gemini-code-assist[bot] commented on code in PR #18917:
URL: https://github.com/apache/tvm/pull/18917#discussion_r2963809939
##########
include/tvm/topi/transform.h:
##########
@@ -904,28 +904,41 @@ inline Tensor strided_slice_with_axes(const Tensor& x,
const ffi::Array<Integer>
std::string slice_mode = "end",
std::string name =
"T_strided_slice_with_axes",
std::string tag = kInjective) {
- const size_t src_tensor_dim = x->shape.size();
- TVM_FFI_ICHECK(axes.size() <= src_tensor_dim);
+ const int64_t src_tensor_dim = static_cast<int64_t>(x->shape.size());
+ TVM_FFI_ICHECK(static_cast<int64_t>(axes.size()) <= src_tensor_dim);
TVM_FFI_ICHECK(axes.size() == begin.size() && axes.size() == end.size() &&
axes.size() == strides.size());
+ // Normalize negative axes
+ ffi::Array<Integer> normalized_axes;
+ for (size_t i = 0; i < axes.size(); ++i) {
+ int64_t axis = axes[i].IntValue();
+ if (axis < 0) {
+ axis += src_tensor_dim;
+ }
+ TVM_FFI_ICHECK(axis >= 0 && axis < src_tensor_dim)
+ << "Axis " << axes[i].IntValue() << " is out of bounds for tensor with
" << src_tensor_dim
+ << " dimensions";
+ normalized_axes.push_back(Integer(axis));
+ }
Review Comment:
While normalizing axes, it's good practice to also check for duplicate axes.
Currently, if an axis is repeated, the slice parameters for that axis will be
silently overwritten, which can lead to unexpected behavior. Other operators
like `transpose` already perform this check. You can use a `std::unordered_set`
to efficiently track seen axes and ensure each axis is specified only once.
```suggestion
ffi::Array<Integer> normalized_axes;
std::unordered_set<int64_t> seen_axes;
for (size_t i = 0; i < axes.size(); ++i) {
int64_t axis = axes[i].IntValue();
if (axis < 0) {
axis += src_tensor_dim;
}
TVM_FFI_ICHECK(axis >= 0 && axis < src_tensor_dim)
<< "Axis " << axes[i].IntValue() << " is out of bounds for tensor
with " << src_tensor_dim
<< " dimensions";
auto [it, inserted] = seen_axes.insert(axis);
TVM_FFI_ICHECK(inserted) << "Axis " << axis << " is specified more than
once.";
normalized_axes.push_back(Integer(axis));
}
```
--
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]
