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?