I am trying to port a complex algorithm from tensorflow to mxnet. It is using the setdiff1d() operator that is not present in mxnet due to the lack of dynamic shape inference. Since I use it in a contrib.while_loop it is not possible to feed it in via let’s say the dataloader.
Any idea how I could implement the same (or very similar) functionality?
For the sake of simplicity lets assume both inputs have unique int values, so I need something like a ‘negative pick’
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)