Hi,
I’m trying to index tensor in the following fashion: x[0, mask]
where x is the input tensor (e.g. (2, 3) shape), while the mask is a boolean array (e.g. (3,) shape). For a better example, here are example values: x = [ [1, 2, 3], [4, 5, 6] ] mask = [True, True, False]
Therefore, when using NumPy as a reference, the results of x[0, mask] should be [1, 2] tensor. However, as TVM output I’m getting something like [2, 2, 1]. Can someone provide me with more information on this topic?
Lastly, for easier reproduction, below is TVM test similar to the example below. FYI, I added it in the following test file: tests/python/relay/test_op_level3.py
Unit test:
def test_adv_index_with_mask(target, dev, executor_kind):
def verify_adv_index(data_shape, index, bool_mask_shape):
dtype = "float32"
inputs = [relay.var("data", relay.TensorType(data_shape, dtype))]
np_data = np.random.uniform(size=data_shape).astype(dtype)
np_indices = np.ones(bool_mask_shape, dtype=bool)
np_indices[-1] = False
inputs.append(relay.var("index_{}".format(index), shape=(), dtype="int"))
inputs.append(relay.var("mask_{}".format(index), relay.TensorType(bool_mask_shape, "bool")))
np_out = np_data[index, np_indices]
np_args = [np_data] + [index, np_indices]
out = relay.op.adv_index(inputs)
func = relay.Function(inputs, out)
op_res = relay.create_executor(executor_kind, device=dev, target=target).evaluate(func)(
*np_args
)
tvm.testing.assert_allclose(op_res.numpy(), np_out, rtol=1e-5)
verify_adv_index((2, 3), 1, (3,))