Add Range op to ONNX, make tvm arange op support negative steps (apac…

…he#6647)
TusharKanekiDey · Oct 15, 2020 · c63ada1 · c63ada1
1 parent f43fbb5
commit c63ada1
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
@@ -1893,6 +1893,19 @@ def _impl_v1(cls, inputs, attr, params):
         return _op.topk(inputs[0], inputs[1], axis=axis)
 
 
+class Range(OnnxOpConverter):
+    """Operator converter for Range"""
+
+    @classmethod
+    def _impl_v1(cls, inputs, attr, params):
+        if len(inputs) != 3:
+            raise ValueError("Expect 3 input only")
+
+        return _op.arange(
+            inputs[0], inputs[1], inputs[2], dtype=infer_type(inputs[0]).checked_type.dtype
+        )
+
+
 class MaxRoiPool(OnnxOpConverter):
     """Operator converter for MaxRoiPool."""