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
(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)
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?