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 ad0a225074 [Fix][Relax][Frontend][ONNX] Fix Scatter with indices
smaller than data (#20187)
ad0a225074 is described below
commit ad0a2250742d26948dd0ec50e980aeb33d5b8712
Author: HuEnwei <[email protected]>
AuthorDate: Fri Aug 28 04:56:58 2026 +0800
[Fix][Relax][Frontend][ONNX] Fix Scatter with indices smaller than data
(#20187)
Fixes: #20182
## Summary
The Relax ONNX frontend imported legal **opset-9/10 Scatter** models
whose
`indices`/`updates` are smaller than `data` (e.g. size-1 broadcast dims)
but
**silently produced wrong numeric output**. The spec iterates over
`indices'`
own shape — `output[idx[:axis] + (indices[idx],) + idx[axis+1:]] =
updates[idx]`
for each entry `idx` — while the frontend forwarded `indices`/`updates`
straight
to `relax.op.scatter_elements`, which only matches that semantics when
`indices.shape == data.shape`.
## Root cause
`Scatter._impl_v9` was a trivial forwarding to
`relax.op.scatter_elements`:
```python
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
axis = attr.get("axis", 0)
return relax.op.scatter_elements(inputs[0], inputs[1], inputs[2],
axis=axis)
```
`scatter_elements` (torch-`scatter_` semantics) writes for every
position of the
broadcast `indices`; ONNX Scatter instead iterates `indices'` own shape.
When
`indices` is smaller than `data`, the two disagree, and the lower-level
op
silently emits wrong values for the size-1-dim (broadcast) cases — e.g.
`data(2,3,4)`, `indices(1,3,1)`, `axis=0` writes updates to the wrong
cells
(`max|diff| = 298` in the minimal repro). A 72-case sweep (3 axes x
indices dims
in {1, dim} x 3 seeds) showed **12/72 wrong**, all with non-axis-1
broadcast dims.
## Fix
When `indices` is statically known and its shape differs from `data`'s,
lower the
per-entry semantics exactly as `scatter_nd` with explicit target
positions: a
constant coordinate grid of `indices'` own shape with the axis column
replaced by
the flattened `indices` values.
```python
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
...
# indices with a dynamic shape: keep the previous lowering.
if not all(isinstance(s, (tirx.IntImm, int)) for s in indices_shape):
return relax.op.scatter_elements(data, indices, updates, axis=axis)
# When indices has data's exact shape, scatter_elements is exact too.
if all(isinstance(s, (tirx.IntImm, int)) for s in data_shape) and list(
indices_shape
) == list(data_shape):
return relax.op.scatter_elements(data, indices, updates, axis=axis)
# per-entry targets: (n_entries, rank) grid of indices' own shape with
the
# axis column replaced by the flattened indices values -> exact
scatter_nd.
rank = len(data_shape)
axis = axis % rank
shape = tuple(int(s) for s in indices_shape)
n_entries = int(_np.prod(shape))
grid = _np.moveaxis(_np.indices(shape), 0, -1).reshape(n_entries, rank)
target = relax.op.where(
relax.const(
_np.broadcast_to(_np.eye(rank, dtype="bool")[axis], (n_entries,
rank)),
"bool",
),
relax.op.reshape(indices, (n_entries, 1)),
relax.const(grid.astype("int64"), "int64"),
)
return relax.op.scatter_nd(data, target, relax.op.reshape(updates,
(n_entries,)))
```
`scatter_nd` has no axis (coordinates are explicit), so negative axes
are
normalized with `axis % rank`. The exact-shape and dynamic-shape cases
keep the
original `scatter_elements` path, so conventional models are unchanged.
This also makes Scatter with mismatched `indices` correct for `axis=0/2`
(which
`test_scatter` previously skipped), so that skip is now restricted to
`ScatterElements` (whose `_impl_v11` still uses `scatter_elements`).
## Validation
Differential test (Relax `from_onnx` + `relax.build` + `VirtualMachine`
vs
onnxruntime) over 72 legal Scatter models: `data(2,3,4)` x `indices`
each dim in
{1, dim} x axes {0,1,2} x 3 seeds, plus the minimal repro, `(16,16,16)+
(8,8,8)` all axes, negative-axis and opset-10 cases. Verified on the
familyfuzz
locked build `262c6d2e0` via runtime monkey-patch (no source files
modified).
| Category | Cases | Before | After |
|---|---|---|---|
| broadcast / smaller `indices` (incl. size-1 dims) | 72 | **12 wrong
output** | match onnxrt, `max\|diff\| = 0` |
| `indices.shape == data.shape` (all axes incl. -1) | 4 | match | match
(no regression) |
| opset 10 broadcast | 2 | wrong | match |
| **Total** | **78** | 12 wrong | **0 wrong** |
In-tree tests added to `tests/python/relax/test_frontend_onnx.py`:
- `test_scatter_broadcast` — 4 axes x 4 broadcast `indices` shapes (16
cases),
`check_correctness` vs onnxruntime with `check_dtypes=True`: 16/16 pass.
- `test_scatter` — the `axis != 1` skip is now restricted to
`ScatterElements`;
Scatter runs for `axis=0/1/2` and passes.
Run:
```bash
pytest tests/python/relax/test_frontend_onnx.py::test_scatter \
tests/python/relax/test_frontend_onnx.py::test_scatter_broadcast
```
## Files changed
- `python/tvm/relax/frontend/onnx/onnx_frontend.py` — `Scatter._impl_v9`
lowers
shape-mismatched (broadcastable) `indices` as exact per-entry
`scatter_nd`.
- `tests/python/relax/test_frontend_onnx.py` — add
`test_scatter_broadcast`;
narrow the `test_scatter` skip to `ScatterElements`.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 72 ++++++++++++++-
tests/python/relax/test_frontend_onnx.py | 117 +++++++++++++++++++++++-
2 files changed, 186 insertions(+), 3 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index a6a2a3152e..d078bdf9f6 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1401,13 +1401,83 @@ class GatherND(OnnxOpConverter):
return relax.op.gather_nd(inputs[0], inputs[1], batch_dims)
+def _shapes_equal(a: list[tirx.Expr] | None, b: list[tirx.Expr] | None) ->
bool:
+ if a is None or b is None or len(a) != len(b):
+ return False
+ for x, y in zip(a, b):
+ x_static = isinstance(x, tirx.IntImm | int)
+ y_static = isinstance(y, tirx.IntImm | int)
+ if x_static and y_static:
+ if int(x) != int(y):
+ return False
+ elif (not x_static) and (not y_static):
+ if not x.same_as(y):
+ return False
+ else:
+ return False
+ return True
+
+
class Scatter(OnnxOpConverter):
"""Convert an onnx Scatter node into an equivalent Relax expression."""
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ indices = inputs[1]
+ updates = inputs[2]
axis = attr.get("axis", 0)
- return relax.op.scatter_elements(inputs[0], inputs[1], inputs[2],
axis=axis)
+
+ indices_shape = indices.ty.shape
+ data_shape = data.ty.shape
+
+ if _shapes_equal(indices_shape, data_shape):
+ return relax.op.scatter_elements(data, indices, updates, axis=axis)
+
+ if indices_shape is None:
+ raise ValueError(
+ "Scatter with `indices` of unknown rank is unsupported, as the
per-entry "
+ "coordinate grid cannot be built"
+ )
+ if not all(isinstance(s, tirx.IntImm | int) for s in indices_shape):
+ raise ValueError(
+ "Scatter with dynamic `indices` whose shape is not provably
equal to "
+ "`data`'s shape is unsupported: the fallback lowering silently
produces "
+ "incorrect results for broadcast size-1 dims"
+ )
+ if data_shape is None:
+ raise ValueError(
+ "Scatter with `data` of unknown rank is unsupported, as the
per-entry "
+ "coordinate grid cannot be built"
+ )
+
+ # ONNX Scatter iterates over indices' own shape: for each entry idx,
+ # output[idx[:axis] + (indices[idx],) + idx[axis+1:]] = updates[idx]
+ # which is exactly scatter_nd with explicit per-entry target positions.
+ # The previous lowering passed a smaller indices (e.g. broadcastable
+ # size-1 dims) directly to scatter_elements, which silently produced
+ # wrong values.
+ rank = len(data_shape)
+ axis = axis % rank
+ # `StructInfo.dtype` is a PrimType, so coerce to a string for
numpy/relax.
+ indices_dtype = str(indices.ty.dtype)
+ shape = tuple(int(s) for s in indices_shape)
+ n_entries = int(_np.prod(shape))
+ # (n_entries, rank) coordinate grid of indices' own shape, C-order over
+ # its entries.
+ grid = _np.moveaxis(_np.indices(shape), 0, -1).reshape(n_entries, rank)
+ # Replace the axis column with the flattened indices values. The grid
is
+ # cast to the indices dtype (ONNX Scatter permits int32 indices) so
that
+ # both branches of `where` share a dtype.
+ target = relax.op.where(
+ relax.const(
+ _np.broadcast_to(_np.eye(rank, dtype="bool")[axis],
(n_entries, rank)),
+ "bool",
+ ),
+ relax.op.reshape(indices, (n_entries, 1)),
+ relax.const(grid.astype(indices_dtype), indices_dtype),
+ )
+ return relax.op.scatter_nd(data, target, relax.op.reshape(updates,
(n_entries,)))
@classmethod
def _impl_v11(cls, bb, inputs, attr, params):
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 4654a4082a..4e31015b10 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1800,8 +1800,11 @@ def test_gather_nd(data_shape, indices_shape,
batch_dims):
@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):
- if axis != 1:
- pytest.skip("The current topi impl is wrong, which only works for
axis=1")
+ if name == "ScatterElements" and axis != 1:
+ pytest.skip(
+ "ScatterElements with indices smaller than data is lowered via
scatter_elements, "
+ "which only works for axis=1"
+ )
input_shape = [16, 16, 16]
indices_shape = [8, 8, 8]
updates_shape = [8, 8, 8]
@@ -1822,6 +1825,116 @@ def test_scatter(axis: int, name: str, opset: int):
check_correctness(model, inputs={"indices": indices}, opset=opset)
[email protected]("indices_dtype", ["int32", "int64"])
+def test_scatter_broadcast(indices_dtype):
+ """Scatter (opset 9/10) whose indices/updates are smaller than data (size-1
+ broadcast dims) must follow the per-entry semantics of the ONNX spec:
+ output[idx[:axis] + (indices[idx],) + idx[axis+1:]] = updates[idx] for each
+ entry idx in indices' own shape. The previous lowering passed such indices
+ directly to scatter_elements, which silently produced wrong values. Both
+ ONNX-permitted indices dtypes (int32/int64) are covered, since the lowered
+ coordinate grid must match the indices dtype."""
+ data_shape = (2, 3, 4)
+ indices_proto = TensorProto.INT32 if indices_dtype == "int32" else
TensorProto.INT64
+ rng = np.random.RandomState(0)
+ for axis in [0, 1, 2, -1]:
+ for indices_shape in [(1, 3, 4), (2, 1, 4), (1, 3, 1), (1, 1, 1)]:
+ graph = helper.make_graph(
+ [
+ helper.make_node(
+ "Scatter", ["data", "indices", "updates"], ["output"],
axis=axis
+ )
+ ],
+ "scatter_broadcast_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT,
data_shape),
+ helper.make_tensor_value_info("indices", indices_proto,
list(indices_shape)),
+ helper.make_tensor_value_info(
+ "updates", TensorProto.FLOAT, list(indices_shape)
+ ),
+ ],
+ outputs=[helper.make_tensor_value_info("output",
TensorProto.FLOAT, data_shape)],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="scatter_broadcast_test",
+ opset_imports=[helper.make_opsetid("", 9)],
+ )
+ inputs = {
+ "data": rng.randn(*data_shape).astype("float32"),
+ "indices": rng.randint(0, data_shape[axis % len(data_shape)],
indices_shape).astype(
+ indices_dtype
+ ),
+ "updates": rng.randn(*indices_shape).astype("float32"),
+ }
+ check_correctness(model, inputs=inputs, opset=9, check_dtypes=True)
+
+
+def test_scatter_dynamic_shape():
+ """Dynamic-shape Scatter: when indices is structurally provably equal to
data
+ (shared symbolic dims) it still lowers via scatter_elements and is correct;
+ a dynamic indices that is *not* provably equal (e.g. a broadcast size-1
dim)
+ must raise instead of silently emitting wrong values."""
+ n = tvm.tirx.Var("N", "int64")
+ rng = np.random.RandomState(1)
+ batch = 5
+
+ # data/indices/updates all share the symbolic batch dim N -> provably
equal,
+ # lowered via scatter_elements, correct per ONNX semantics.
+ shape_dict = {"data": [n, 3, 4], "indices": [n, 3, 4], "updates": [n, 3,
4]}
+ graph = helper.make_graph(
+ [helper.make_node("Scatter", ["data", "indices", "updates"],
["output"], axis=0)],
+ "scatter_dynamic_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT, ["N", 3,
4]),
+ helper.make_tensor_value_info("indices", TensorProto.INT64, ["N",
3, 4]),
+ helper.make_tensor_value_info("updates", TensorProto.FLOAT, ["N",
3, 4]),
+ ],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT,
["N", 3, 4])],
+ )
+ model = helper.make_model(
+ graph, producer_name="scatter_dynamic_test",
opset_imports=[helper.make_opsetid("", 9)]
+ )
+ model.ir_version = 8
+ data = rng.randn(batch, 3, 4).astype("float32")
+ indices = rng.randint(0, batch, size=(batch, 3, 4)).astype("int64")
+ updates = rng.randn(batch, 3, 4).astype("float32")
+ ort_output = onnxruntime.InferenceSession(
+ model.SerializeToString(), providers=["CPUExecutionProvider"]
+ ).run(None, {"data": data, "indices": indices, "updates": updates})[0]
+ tvm_model = from_onnx(model, shape_dict=shape_dict, opset=9)
+ tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
+ tvm_model = relax.transform.LegalizeOps()(tvm_model)
+ with tvm.transform.PassContext(opt_level=3):
+ ex = tvm.compile(tvm_model, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ vm.set_input("main", data, indices, updates)
+ vm.invoke_stateful("main")
+ tvm_output = vm.get_outputs("main")
+ np.testing.assert_allclose(ort_output, tvm_output.numpy(), rtol=1e-6,
atol=1e-6)
+
+ # A dynamic indices that is not provably equal to data (size-1 broadcast
dim)
+ # must raise rather than silently produce wrong values.
+ shape_dict_bad = {"data": [n, 3, 4], "indices": [n, 1, 4], "updates": [n,
1, 4]}
+ graph_bad = helper.make_graph(
+ [helper.make_node("Scatter", ["data", "indices", "updates"],
["output"], axis=0)],
+ "scatter_dynamic_bad_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT, ["N", 3,
4]),
+ helper.make_tensor_value_info("indices", TensorProto.INT64, ["N",
1, 4]),
+ helper.make_tensor_value_info("updates", TensorProto.FLOAT, ["N",
1, 4]),
+ ],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT,
["N", 3, 4])],
+ )
+ model_bad = helper.make_model(
+ graph_bad,
+ producer_name="scatter_dynamic_bad_test",
+ opset_imports=[helper.make_opsetid("", 9)],
+ )
+ with pytest.raises(ValueError, match="dynamic"):
+ from_onnx(model_bad, shape_dict=shape_dict_bad, opset=9)
+
+
@pytest.mark.parametrize(
"reduction, opset, data, indices, updates",
[
