This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 3aa0eb16f7 [Fix][Relax] Normalize negative indices in
Gather/Scatter/OneHot Ops (#20219)
3aa0eb16f7 is described below
commit 3aa0eb16f70e42a7ed716b2e5a0642eda342b3fe
Author: Kryptonite <[email protected]>
AuthorDate: Sat Aug 29 08:02:47 2026 +0300
[Fix][Relax] Normalize negative indices in Gather/Scatter/OneHot Ops
(#20219)
GatherElements, GatherND, ScatterND, and OneHot in the ONNX frontend
forwarded negative indices as-is instead of normalizing them (idx +
axis_length).
Fixes #20217.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 100 +++++++++++++++-----
tests/python/relax/test_frontend_onnx.py | 118 ++++++++++++++++++++++++
2 files changed, 196 insertions(+), 22 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 0b2e5063c8..8a40e4a612 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1329,6 +1329,48 @@ class Cast(OnnxOpConverter):
return relax.op.astype(inputs[0], to_type)
+def _normalize_negative_indices(bb, indices, axis_extent):
+ """Map negative indices to idx + axis_extent; skip unsigned indices.
axis_extent
+ must be broadcastable against indices and share its dtype.
+ """
+ indices_dtype = indices.ty.dtype.dtype
+ if indices_dtype.startswith("uint"):
+ return indices
+ return bb.normalize(
+ relax.op.where(
+ relax.op.less(indices, relax.const(0, indices_dtype)),
+ relax.op.add(indices, axis_extent),
+ indices,
+ )
+ )
+
+
+def _axis_extent(bb, data, axis, indices_dtype):
+ """Compute data.shape[axis] at runtime, cast to indices_dtype."""
+ data_shape_tensor =
bb.normalize(relax.op.shape_to_tensor(relax.op.shape_of(data)))
+ axis_extent = bb.normalize(
+ relax.op.take(data_shape_tensor, relax.const(axis, "int64"), axis=0,
mode="wrap")
+ )
+ if indices_dtype != "int64":
+ axis_extent = bb.normalize(relax.op.astype(axis_extent, indices_dtype))
+ return axis_extent
+
+
+def _coord_axis_extents(bb, data, indices, start_axis):
+ """Compute data.shape[start_axis + j] for each coordinate slot j, cast to
indices_dtype."""
+ indices_dtype = indices.ty.dtype.dtype
+ coord_len = indices.ty.shape[-1]
+ data_shape_tensor =
bb.normalize(relax.op.shape_to_tensor(relax.op.shape_of(data)))
+ axis_extent = bb.normalize(
+ relax.op.strided_slice(
+ data_shape_tensor, axes=[0], begin=[start_axis], end=[start_axis +
coord_len]
+ )
+ )
+ if indices_dtype != "int64":
+ axis_extent = bb.normalize(relax.op.astype(axis_extent, indices_dtype))
+ return axis_extent
+
+
class Gather(OnnxOpConverter):
"""Convert an onnx Gather node into an equivalent Relax expression."""
@@ -1365,23 +1407,8 @@ class Gather(OnnxOpConverter):
indices = bb.normalize(relax.op.shape_to_tensor(indices))
indices_dtype = indices.ty.dtype.dtype
- if not indices_dtype.startswith("uint"):
- data_shape = bb.normalize(relax.op.shape_of(data))
- data_shape_tensor =
bb.normalize(relax.op.shape_to_tensor(data_shape))
- axis_extent = bb.normalize(
- relax.op.take(data_shape_tensor, relax.const(axis, "int64"),
axis=0, mode="wrap")
- )
-
- if indices_dtype != "int64":
- axis_extent = bb.normalize(relax.op.astype(axis_extent,
indices_dtype))
-
- indices = bb.normalize(
- relax.op.where(
- relax.op.less(indices, relax.const(0, indices_dtype)),
- relax.op.add(indices, axis_extent),
- indices,
- )
- )
+ axis_extent = _axis_extent(bb, data, axis, indices_dtype)
+ indices = _normalize_negative_indices(bb, indices, axis_extent)
return relax.op.take(data, indices, axis)
@@ -1391,8 +1418,15 @@ class GatherElements(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ indices = inputs[1]
axis = attr.get("axis", 0)
- return relax.op.gather_elements(inputs[0], inputs[1], axis)
+
+ indices_dtype = indices.ty.dtype.dtype
+ axis_extent = _axis_extent(bb, data, axis, indices_dtype)
+ indices = _normalize_negative_indices(bb, indices, axis_extent)
+
+ return relax.op.gather_elements(data, indices, axis)
class GatherND(OnnxOpConverter):
@@ -1400,8 +1434,13 @@ class GatherND(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ indices = inputs[1]
batch_dims = attr.get("batch_dims", 0)
- return relax.op.gather_nd(inputs[0], inputs[1], batch_dims)
+
+ axis_extent = _coord_axis_extents(bb, data, indices, batch_dims)
+ indices = _normalize_negative_indices(bb, indices, axis_extent)
+ return relax.op.gather_nd(data, indices, batch_dims)
def _shapes_equal(a: list[tirx.Expr] | None, b: list[tirx.Expr] | None) ->
bool:
@@ -1531,19 +1570,31 @@ class ScatterND(OnnxOpConverter):
def _reduction_check(attr, valid_reductions: list[str]):
return _get_onnx_reduction(attr, valid_reductions)
+ @staticmethod
+ def _normalize_indices(bb, data, indices):
+ # ScatterND has no batch_dims: coordinate slot j always indexes data
axis j.
+ axis_extent = _coord_axis_extents(bb, data, indices, start_axis=0)
+ return _normalize_negative_indices(bb, indices, axis_extent)
+
@classmethod
def _impl_v11(cls, bb, inputs, attr, params):
- return relax.op.scatter_nd(inputs[0], inputs[1], inputs[2])
+ data, indices, updates = inputs
+ indices = cls._normalize_indices(bb, data, indices)
+ return relax.op.scatter_nd(data, indices, updates)
@classmethod
def _impl_v16(cls, bb, inputs, attr, params):
+ data, indices, updates = inputs
reduction = cls._reduction_check(attr, ["update", "add", "mul"])
- return relax.op.scatter_nd(inputs[0], inputs[1], inputs[2], reduction)
+ indices = cls._normalize_indices(bb, data, indices)
+ return relax.op.scatter_nd(data, indices, updates, reduction)
@classmethod
def _impl_v18(cls, bb, inputs, attr, params):
+ data, indices, updates = inputs
reduction = cls._reduction_check(attr, ["update", "add", "mul", "min",
"max"])
- return relax.op.scatter_nd(inputs[0], inputs[1], inputs[2], reduction)
+ indices = cls._normalize_indices(bb, data, indices)
+ return relax.op.scatter_nd(data, indices, updates, reduction)
class Compress(OnnxOpConverter):
@@ -5077,6 +5128,11 @@ class OneHot(OnnxOpConverter):
relax.prim_value(off_value),
relax.prim_value(on_value),
)
+
+ indices_dtype = indices.ty.dtype.dtype
+ axis_extent = relax.const(depth, indices_dtype)
+ indices = _normalize_negative_indices(bb, indices, axis_extent)
+
return relax.op.one_hot(indices, on_value, off_value, depth, axis)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 00c34b775b..a2081adab9 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1816,6 +1816,37 @@ def test_gather_elements(data_shape, indices_shape,
axis):
check_correctness(model, inputs=input_values)
[email protected](
+ "data_shape, indices_shape, axis",
+ [
+ ([3, 4, 5], [1, 4, 5], 0),
+ ([3, 4, 5], [3, 2, 5], 1),
+ ([3, 4, 5], [3, 4, 2], 2),
+ ],
+)
+def test_gather_elements_negative_indices(data_shape, indices_shape, axis):
+ gather_elements_node = helper.make_node("GatherElements", ["data",
"indices"], ["y"], axis=axis)
+
+ graph = helper.make_graph(
+ [gather_elements_node],
+ "gather_elements_negative_indices_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT,
data_shape),
+ helper.make_tensor_value_info("indices", TensorProto.INT64,
indices_shape),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
indices_shape)],
+ )
+
+ model = helper.make_model(graph,
producer_name="gather_elements_negative_indices_test")
+ input_values = {
+ "data": np.random.randn(*data_shape).astype("float32"),
+ "indices": np.random.randint(-data_shape[axis], data_shape[axis],
indices_shape).astype(
+ "int64"
+ ),
+ }
+ check_correctness(model, inputs=input_values)
+
+
@pytest.mark.parametrize(
"data_shape, indices_shape, batch_dims",
[
@@ -1849,6 +1880,39 @@ def test_gather_nd(data_shape, indices_shape,
batch_dims):
check_correctness(model, inputs=input_values)
[email protected](
+ "data_shape, indices_shape, batch_dims",
+ [
+ ([2, 2], [2, 2], 0),
+ ([2, 2], [2, 1], 0),
+ ([2, 2, 2], [1], 0),
+ ([2, 2, 2], [2, 2], 0),
+ ([2, 2, 2], [2, 1, 2], 0),
+ ([2, 2, 2], [2, 2], 1),
+ ([2, 2, 2], [2, 1], 1),
+ ],
+)
+def test_gather_nd_negative_indices(data_shape, indices_shape, batch_dims):
+ gather_nd_node = helper.make_node("GatherND", ["data", "indices"], ["y"],
batch_dims=batch_dims)
+
+ graph = helper.make_graph(
+ [gather_nd_node],
+ "gather_nd_negative_indices_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT,
data_shape),
+ helper.make_tensor_value_info("indices", TensorProto.INT64,
indices_shape),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)],
+ )
+
+ model = helper.make_model(graph,
producer_name="gather_nd_negative_indices_test")
+ input_values = {
+ "data": np.random.randn(*data_shape).astype("float32"),
+ "indices": np.random.randint(-2, 2, indices_shape).astype("int64"),
+ }
+ check_correctness(model, inputs=input_values)
+
+
@pytest.mark.parametrize("axis", [0, 1, 2])
@pytest.mark.parametrize(("name", "opset"), [("Scatter", 10),
("ScatterElements", 11)])
def test_scatter(axis: int, name: str, opset: int):
@@ -2119,6 +2183,38 @@ def test_scatter_nd(reduction):
verify_scatter_nd([10], [5, 1], [5])
[email protected]("reduction", ["none", "add", "mul"])
+def test_scatter_nd_negative_indices(reduction):
+ def verify_scatter_nd_negative_indices(data_shape, indices_shape,
updates_shape):
+ scatter_nd_node = helper.make_node(
+ "ScatterND",
+ ["data", "indices", "updates"],
+ ["output"],
+ reduction=reduction,
+ )
+
+ graph = helper.make_graph(
+ [scatter_nd_node],
+ "scatter_nd_negative_indices_test",
+ inputs=[
+ 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("updates", TensorProto.FLOAT,
updates_shape),
+ ],
+ outputs=[helper.make_tensor_value_info("output",
TensorProto.FLOAT, data_shape)],
+ )
+
+ model = helper.make_model(graph,
producer_name="scatter_nd_negative_indices_test")
+
+ indices = np.random.randint(-data_shape[0], data_shape[0],
indices_shape)
+ check_correctness(model, inputs={"indices": indices}, opset=16)
+
+ verify_scatter_nd_negative_indices([8], [4, 1], [4])
+ verify_scatter_nd_negative_indices([4, 4, 4], [2, 1], [2, 4, 4])
+ verify_scatter_nd_negative_indices([4, 5, 6], [2, 3, 2], [2, 3, 6])
+ verify_scatter_nd_negative_indices([10], [5, 1], [5])
+
+
def test_compress():
def verify_compress(
tensor_shape: list[int],
@@ -9859,6 +9955,28 @@ def test_onehot():
check_correctness(model, inputs=values)
+def test_onehot_negative_indices():
+ one_hot_node = helper.make_node("OneHot", ["indices", "depth", "values"],
["y"], axis=1)
+ graph = helper.make_graph(
+ [one_hot_node],
+ "one_hot_negative_indices_test",
+ inputs=[
+ helper.make_tensor_value_info("indices", TensorProto.INT64, [2,
2]),
+ ],
+ initializer=[
+ helper.make_tensor("depth", TensorProto.INT64, [], [10]),
+ helper.make_tensor("values", TensorProto.FLOAT, [2], [3, 1]),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 10,
2])],
+ )
+
+ model = helper.make_model(graph,
producer_name="one_hot_negative_indices_test")
+ values = {
+ "indices": np.array([[-1, -10], [-2, 4]], dtype="int64"),
+ }
+ check_correctness(model, inputs=values)
+
+
@pytest.mark.parametrize("axis", [None, 0, 1, -1])
@pytest.mark.parametrize("sorted", [0, 1])
@pytest.mark.parametrize("num_outputs", [1, 2, 3, 4])
