[siyiweigeHEW]

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:

@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.

@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:

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.

[tlopex]

This hard-codes the coordinate grid as int64, but ONNX Scatter permits int32 indices as well. For a valid int32 model, the two branches of relax.op.where have different dtypes, causing import to fail. Please use the indices dtype and add an int32 test. Also, dynamic shapes still fall back to the old lowering that can silently produce incorrect results.

[siyiweigeHEW]

Hi @tlopex, all requested changes have been implemented:

1. Dtype Match: Dynamically bound the coordinate grid dtype to indices.dtype (supporting both int32 and int64).
2. Dynamic Shape Safety: Added an explicit check/ValueError for dynamic shapes that cannot be proven equal, preventing fallback to silent incorrect lowering.
3. Tests: Added int32/int64 test coverage (test_scatter_broadcast) and dynamic shape boundary test (test_scatter_dynamic_shape) with ir_version=8 for ONNXRuntime compatibility.

All CI checks have passed. Could you please re-review when you have time? Thanks!

[tlopex]

Thanks LGTM