tqchen commented on a change in pull request #9023:
URL: https://github.com/apache/tvm/pull/9023#discussion_r709643444



##########
File path: src/te/tensor.cc
##########
@@ -39,15 +39,26 @@ IterVar reduce_axis(Range dom, std::string name) { return 
IterVar(dom, Var(name)
 Var var(std::string name_hint, DataType t) { return Var(name_hint, t); }
 
 // Tensor
-PrimExpr Tensor::operator()(Array<Var> indices) const {
+PrimExpr Tensor::operator()(Array<Var> indices, bool support_negative_indices 
= false) const {
   Array<PrimExpr> arr(indices.begin(), indices.end());
-  return operator()(arr);
+  return operator()(arr, support_negative_indices);
 }
 
-PrimExpr Tensor::operator()(Array<PrimExpr> indices) const {
-  if (ndim() != 0) {
-    ICHECK_EQ(ndim(), indices.size()) << "Tensor dimension mismatch in read "
-                                      << "ndim = " << ndim() << ", 
indices.size=" << indices.size();
+PrimExpr Tensor::operator()(Array<PrimExpr> indices, bool 
support_negative_indices = false) const {
+  Array<PrimExpr> shape = (*this)->shape;
+
+  if (shape.size() != 0) {
+    ICHECK_EQ(shape.size(), indices.size())
+        << "Tensor dimension mismatch in read "
+        << "ndim = " << ndim() << ", indices.size=" << indices.size();
+  }
+
+  if (support_negative_indices) {
+    for (size_t i = 0; i < shape.size(); i++) {
+      PrimExpr new_index = if_then_else(indices[i] < 
make_const(indices[i]->dtype, 0),

Review comment:
       you want to use select here, because it does not involves a memory 
operation(thus can evaluate both sides)




-- 
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: commits-unsubscr...@tvm.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to