The GitHub Actions job "Lint" on tvm.git/main has succeeded. Run started by GitHub user tlopex (triggered by tlopex).
Head commit for run: ad0a2250742d26948dd0ec50e980aeb33d5b8712 / HuEnwei <[email protected]> [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`. Report URL: https://github.com/apache/tvm/actions/runs/33115801539 With regards, GitHub Actions via GitBox --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
