siyiweigeHEW commented on code in PR #20187:
URL: https://github.com/apache/tvm/pull/20187#discussion_r3875516801
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1406,8 +1406,48 @@ class Scatter(OnnxOpConverter):
@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
+
+ # 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 matches ONNX's
+ # per-entry semantics and is the direct lowering.
+ 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)
+
+ # 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
+ 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.
+ 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"),
Review Comment:
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!
--
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]
