 # 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

``````[18. 19. 20. 21. 22. 23.]