I don’t think there is anyway for MXNet to produce a hybridizable computational graph with variable output shape. One option is to have the operation return a mask array with dimension of x (for
setdiff1d(x, y)) and then you can take the mask and do other operations on top (calculate number, run while-loop, etc.).
To get the mask, you can convert x and y to one-hot representation (I assume you know an upper bound for values in x and y), create a negative mask from one-hot encoded y values, multiply this mask by x, and lastly create a positive mask from masked one-hot x array. Here is an implementation:
x_mask = (x.one_hot(max_int) * (1 - y.one_hot(max_int).sum(axis=0, keepdims=True))).sum(axis=1)