Tensor advanced indexing with boolean mask


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)(
        tvm.testing.assert_allclose(op_res.numpy(), np_out, rtol=1e-5)

    verify_adv_index((2, 3), 1, (3,))