[GitHub] [incubator-tvm] mbrookhart commented on a change in pull request #6647: Add Range op to ONNX, make tvm arange shape_func support negative steps

[GitBox <gi...@apache.org>]

mbrookhart commented on a change in pull request #6647:
URL: https://github.com/apache/incubator-tvm/pull/6647#discussion_r501804197



##########
File path: python/tvm/relay/op/_transform.py
##########
@@ -123,7 +123,10 @@ def compute_scatter_add(attrs, inputs, output_type):
 @script
 def _arange_shape_func(start, stop, step):
     out = output_tensor((1,), "int64")
-    out[0] = int64(ceil_div((int64(stop[0]) - int64(start[0])), int64(step[0])))
+    if step[0] < 0:
+        out[0] = int64(ceil_div((int64(start[0]) - int64(stop[0])), int64(-step[0])))
+    else:
+        out[0] = int64(ceil_div((int64(stop[0]) - int64(start[0])), int64(step[0])))

Review comment:
       It looks like there are already tests with negative steps in test_op_level3.py, they just weren't hitting the shape_func




----------------------------------------------------------------
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.

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