gemini-code-assist[bot] commented on code in PR #19967:
URL: https://github.com/apache/tvm/pull/19967#discussion_r3543811248
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1081,17 +1081,21 @@ def _impl_v13(cls, bb, inputs, attr, params):
output = _np.take(data.data.numpy(), indices.data.numpy(),
axis=axis)
return relax.const(output, output.dtype)
- # If input is a shape expression, take a value from that shape and
return it as a constant.
+ # If input is a shape expression, take a value from that shape. A
constant
+ # index resolves to a single dimension that we return as a PrimValue
to keep
+ # shape-specialized handling in downstream shape-construction
patterns. A
+ # dynamic index materializes the shape as an int64 tensor and gathers
from it
+ # at runtime, matching how Reshape handles ShapeExpr inputs.
if isinstance(data, relax.ShapeExpr):
- assert isinstance(indices, relax.Constant), (
- "Only constant indices supported for shape gather."
- )
- np_index = indices.data.numpy()
- if len(np_index.shape) == 1:
- np_index = np_index[0]
- np_index = int(np_index)
- shape_val = data[np_index]
- return relax.PrimValue(shape_val)
+ if isinstance(indices, relax.Constant):
+ np_index = indices.data.numpy()
+ if len(np_index.shape) == 1:
+ np_index = np_index[0]
+ np_index = int(np_index)
+ shape_val = data[np_index]
+ return relax.PrimValue(shape_val)
Review Comment:
There is a correctness issue in the constant-index fast path when `indices`
is a constant of size greater than 1 (e.g., `[0, 2]`).
Currently, if `indices` is a constant of size > 1, `len(np_index.shape) ==
1` is `True`, so `np_index = np_index[0]` extracts only the first element
(e.g., `0`), and the function returns a single scalar
`relax.PrimValue(data[0])`. This silently discards the other indices and
returns a scalar instead of a 1D tensor of the gathered dimensions.
To fix this, the fast path should only be taken if `indices` is a constant
of size 1 (i.e., a scalar or a 1D tensor of size 1). If the constant has a size
greater than 1, it should fall through to the dynamic path where `data` is
converted to a tensor and `relax.op.take` is used, which correctly handles
multi-element indices.
```suggestion
if isinstance(indices, relax.Constant) and
indices.data.numpy().size == 1:
np_index = indices.data.numpy()
if len(np_index.shape) == 1:
np_index = np_index[0]
np_index = int(np_index)
shape_val = data[np_index]
return relax.PrimValue(shape_val)
```
--
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]
