gemini-code-assist[bot] commented on code in PR #19968:
URL: https://github.com/apache/tvm/pull/19968#discussion_r3544110174
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1182,17 +1182,22 @@ 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
single
+ # constant index resolves to one dimension that we return as a
PrimValue to
+ # keep shape-specialized handling in downstream shape-construction
patterns.
+ # Any other index (dynamic, or a constant selecting multiple
dimensions)
+ # materializes the shape as an int64 tensor and gathers from it at
runtime,
+ # reusing the negative-index normalization below.
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.prim_value(shape_val)
+ 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.prim_value(shape_val)
Review Comment:
We can improve this check to be more robust and efficient:
1. **Avoid redundant calls**: `indices.data.numpy()` is currently called
multiple times, which creates a new NumPy array copy/view on each call.
2. **Prevent potential `TypeError`**: If `indices` is a higher-dimensional
constant of size 1 (e.g., `[[3]]`), `int(np_index)` will raise a `TypeError` in
Python/NumPy. Restricting the fast path to `ndim <= 1` prevents this.
3. **Simplify indexing**: Using `.item()` is a standard and clean way to
extract the scalar value of a single-element array, regardless of whether it is
0-D or 1-D.
```suggestion
if isinstance(indices, relax.Constant):
np_indices = indices.data.numpy()
if np_indices.ndim <= 1 and np_indices.size == 1:
np_index = int(np_indices.item())
shape_val = data[np_index]
return relax.prim_value(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]
