gemini-code-assist[bot] commented on code in PR #19525:
URL: https://github.com/apache/tvm/pull/19525#discussion_r3212633539
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1106,6 +1106,25 @@ def _impl_v13(cls, bb, inputs, attr, params):
shape_val = data[np_index]
return relax.PrimValue(shape_val)
+ data_shape = bb.normalize(relax.op.shape_of(data))
+ data_shape_tensor = bb.normalize(relax.op.shape_to_tensor(data_shape))
+ axis_extent = bb.normalize(
+ relax.op.take(data_shape_tensor, relax.const(axis, "int64"),
axis=0, mode="wrap")
+ )
+
+ indices_dtype = indices.struct_info.dtype
+ if not indices_dtype.startswith("uint"):
+ if indices_dtype !="int64":
+ axis_extent = bb.normalize(relax.op.astype(axis_extent,
indices_dtype))
+
+ indices = bb.normalize(
+ relax.op.where(
+ relax.op.less(indices, relax.const(0, indices_dtype)),
+ relax.op.add(indices, axis_extent),
+ indices,
+ )
+ )
Review Comment:
The calculation of `axis_extent` and its supporting shape tensors should be
moved inside the `if not indices_dtype.startswith("uint"):` block. Currently,
these operations are emitted into the Relax graph even when the indices are
unsigned (where normalization is unnecessary), leading to redundant
computations at runtime. Additionally, there is a minor style issue with
missing spaces around the `!=` operator.
```python
indices_dtype = indices.struct_info.dtype
if not indices_dtype.startswith("uint"):
data_shape = bb.normalize(relax.op.shape_of(data))
data_shape_tensor =
bb.normalize(relax.op.shape_to_tensor(data_shape))
axis_extent = bb.normalize(
relax.op.take(data_shape_tensor, relax.const(axis, "int64"),
axis=0, mode="wrap")
)
if indices_dtype != "int64":
axis_extent = bb.normalize(relax.op.astype(axis_extent,
indices_dtype))
indices = bb.normalize(
relax.op.where(
relax.op.less(indices, relax.const(0, indices_dtype)),
relax.op.add(indices, axis_extent),
indices,
)
)
```
--
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]
