Get nd.array indices from condition

Hi,

I’m trying to get same results with NDArray as I get with Numpy when using conditions. I’m interested in getting the indices that are true for a given condition and use these indices to slice another matrix. Example:

import numpy as np

a = np.arange(0, 12).reshape(3, 4)
b = np.arange(12, 24).reshape(3, 4)

idx_np = a > 5
ans1 = b[idx_np]

# or

rows, cols = np.where(a > 5)
ans2 = b[(rows, cols)]

# True
assert np.allclose(ans1, ans2)   

So I am interested in idx_np or (rows, cols). Both can work to do what I want. However, if I do something similar with NDArrays, it does not work:

from mxnet import nd

a = nd.arange(0, 12).reshape(3, 4)
b = nd.arange(12, 24).reshape(3, 4)

idx_nd = a > 5

ans3 = b[idx_nd]

# ValueError: operands could not be broadcast together with shapes (6,) (3,4,4) 
assert np.all(ans3.tonumpy(), ans1)  

Since idx_nd is a 3x4 matrix of 0 and 1 of floats, and not a matrix of booleans like idx_np, ans3 gives me back a 3D matrix with copies of rows, instead of a single array with elements from the intended index.

What is the correct way to do index slicing by condition with NDArray?

Hi @blakec,

Check out mx.nd.contrib.boolean_mask for this. You need to flatten your data to a single dimension before using this operator which is a little different from numpy (where the flatten is done afterwards).

from mxnet import nd

a = nd.arange(0, 12).reshape(3, 4)
b = nd.arange(12, 24).reshape(3, 4)

a = a.reshape(-1)
b = b.reshape(-1)
idx = a > 5

out = mx.nd.contrib.boolean_mask(b, idx)
print(out)
[18. 19. 20. 21. 22. 23.]
<NDArray 6 @cpu(0)>
2 Likes

You’re the best! Thank you so much.

1 Like