gemini-code-assist[bot] commented on code in PR #19881:
URL: https://github.com/apache/tvm/pull/19881#discussion_r3465643192
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -3387,11 +3397,37 @@ def _convert_stablehlo_dynamic_update_slice(self, op):
update_indices[axis] += start
update_indices = np.moveaxis(update_indices, 0, -1)
- operand = self.get_tensor_expr(operand_tensor)
- update = self.get_tensor_expr(update_tensor)
indices = self.bb.normalize(relax.const(update_indices, dtype="int64"))
return self.bb.normalize(relax.op.scatter_nd(operand, indices, update,
"update"))
+ def _build_dynamic_update_slice_indices(self, start_tensors,
operand_shape, update_shape, rank):
+ """Build the scatter_nd index grid for runtime DYNAMIC_UPDATE_SLICE
starts.
+
+ Returns an int64 tensor of shape ``(*update_shape, rank)`` where axis
``a``
+ holds ``arange(update_shape[a]) + clamp(start[a], 0, operand_dim -
update_dim)``,
+ broadcast over the other axes (StableHLO clamps out-of-range starts).
+ """
+ axis_indices = []
+ for axis in range(rank):
+ start_expr = self.bb.normalize(
+ relax.op.astype(self.get_tensor_expr(start_tensors[axis]),
"int64")
+ )
+ max_start = operand_shape[axis] - update_shape[axis]
+ start_expr = relax.op.maximum(start_expr, relax.const(0, "int64"))
+ start_expr = relax.op.minimum(start_expr, relax.const(max_start,
"int64"))
+
+ base = relax.op.arange(update_shape[axis], dtype="int64")
Review Comment:
Calling `relax.op.arange` with only one positional argument
(`update_shape[axis]`) and a keyword argument `dtype` can lead to type
mismatches or signature errors in TVM Relax. Specifically, if `dtype` is
`"int64"`, the implicitly converted `stop` argument (from Python `int`) may
default to `"int32"`, causing a type mismatch with the `"int64"` `start` and
`step` constants.
To ensure robustness and prevent potential compilation or runtime errors, it
is highly recommended to explicitly pass `start`, `stop`, and `step` as
`relax.const` expressions with matching dtypes.
```suggestion
base = relax.op.arange(
relax.const(0, "int64"),
relax.const(update_shape[axis], "int64"),
relax.const(1, "int64"),
)
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -3369,10 +3375,14 @@ def _convert_stablehlo_dynamic_update_slice(self, op):
"and start-index ranks to match"
)
+ operand = self.get_tensor_expr(operand_tensor)
+ update = self.get_tensor_expr(update_tensor)
Review Comment:
Since both `operand_shape` and `update_shape` are statically known, we
should validate at compile time that the update slice dimensions do not exceed
the operand dimensions (i.e., `update_shape[axis] <= operand_shape[axis]` for
all axes). This prevents invalid models from being imported and provides a
clear, early error message instead of failing later during compilation or
runtime.
```suggestion
for dim, size in zip(operand_shape, update_shape):
if size > dim:
raise ValueError(
"STABLEHLO_DYNAMIC_UPDATE_SLICE update shape must be
smaller than "
"or equal to operand shape for all dimensions."
)
operand = self.get_tensor_expr(operand_tensor)
update = self.get_tensor_expr(update_tensor)
```
--
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]
