diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index b2d53198f643..223926d07c87 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -1526,9 +1526,8 @@ inline Tensor gather(const Tensor& data, int axis, const Tensor& indices, } TVM_FFI_ICHECK_GE(axis, 0); TVM_FFI_ICHECK_LT(axis, ndim_d); - if (indices->shape[axis].as()) { - size_t indices_dim_i = static_cast(GetConstInt(indices->shape[axis])); - TVM_FFI_ICHECK_GE(indices_dim_i, 1); + if (const auto* indices_dim = indices->shape[axis].as()) { + TVM_FFI_ICHECK_GE(indices_dim->value, 0); } // Index tensors are validated by integer element kind; vector lane encoding is irrelevant here. PrimType indices_ty = indices->dtype; diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 7f540a53f6f9..96f8fb5659c4 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -2249,6 +2249,40 @@ def test_gather_elements(data_shape, indices_shape, axis): check_correctness(model, inputs=input_values) +@pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="LLVM is required") +@pytest.mark.parametrize( + "data_shape,indices_shape,axis", + [ + ((2, 3), (0, 3), 0), + ((2, 3), (2, 0), -1), + ((0, 3), (0, 3), 0), + ((0, 3), (0, 2), 1), + ((2, 3), (1, 3), 0), + ], +) +def test_gather_elements_with_empty_indices(data_shape, indices_shape, axis): + graph = helper.make_graph( + [helper.make_node("GatherElements", ["data", "indices"], ["y"], axis=axis)], + "gather_elements_with_empty_indices", + [ + helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape), + helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, indices_shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + onnx.checker.check_model(model) + # Give the importer concrete input shapes to isolate TOPI lowering. + mod = from_onnx(model, shape_dict={"data": data_shape, "indices": indices_shape}) + vm = relax.VirtualMachine(tvm.compile(mod, target="llvm"), tvm_tirx.cpu()) + + data = np.arange(np.prod(data_shape), dtype=np.float32).reshape(data_shape) + indices = np.zeros(indices_shape, dtype=np.int64) + actual = vm["main"](tvm_tirx.runtime.tensor(data), tvm_tirx.runtime.tensor(indices)).numpy() + expected = np.take_along_axis(data, indices, axis) + np.testing.assert_array_equal(actual, expected, strict=True) + + @pytest.mark.parametrize( "data_shape, indices_shape, axis", [