tlopex commented on code in PR #20187:
URL: https://github.com/apache/tvm/pull/20187#discussion_r3858885841
##########
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:
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.
--
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]
