siyiweigeHEW opened a new pull request, #20187:
URL: https://github.com/apache/tvm/pull/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`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
