Update docs

docs/api/python/topi.rst

@@ -67,6 +67,7 @@ List of operators
    topi.not_equal
    topi.greater_equal
    topi.less_equal
+   topi.arange
    topi.image.resize
 
 
@@ -123,6 +124,7 @@ topi
 .. autofunction:: topi.power
 .. autofunction:: topi.greater
 .. autofunction:: topi.less
+.. autofunction:: topi.arange
 
 topi.nn
 ~~~~~~~

docs/langref/relay_op.rst

@@ -95,6 +95,7 @@ This level enables additional math and transform operators.
    tvm.relay.full_like
    tvm.relay.cast
    tvm.relay.split
+   tvm.relay.arange
 
 
 **Level 4: Broadcast and Reductions**
@@ -216,6 +217,7 @@ Level 3 Definitions
 .. autofunction:: tvm.relay.full_like
 .. autofunction:: tvm.relay.cast
 .. autofunction:: tvm.relay.split
+.. autofunction:: tvm.relay.arange
 
 
 Level 4 Definitions

python/tvm/relay/op/transform.py

@@ -250,6 +250,39 @@ def full_like(data, fill_value):
 
 
 def arange(stop, start=None, step=1, dtype="float32"):
+    """Return evenly spaced values within a given interval.
+
+    Warning: Undefined behavior when dtype is incompatible with start/stop/step.
+    It could lead to different results compared to numpy, MXNet, pytorch, etc.
+
+    Parameters
+    ----------
+    stop : tvm.Expr
+        Stop of interval. The interval does not include this value.
+
+    start : tvm.Expr, optional
+        Start of interval. The interval includes this value. The default start
+        value is 0.
+
+    step : tvm.Expr, optional
+        Spacing between values. The default step size is 1.
+
+    dtype : str, optional
+        The target data type.
+
+    Returns
+    -------
+    result : relay.Expr
+        The resulting tensor.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        relay.arange(5) = [0, 1, 2, 3, 4]
+        relay.arange(1, 5) = [1, 2, 3, 4]
+        relay.arange(1, 5, 1.5) = [1, 2.5, 4]
+    """
     if start is None:
         start = 0
     else:

src/relay/op/tensor/transform.cc

@@ -891,7 +891,6 @@ bool ArangeRel(const Array<Type>& types,
   const ArangeAttrs* param = attrs.as<ArangeAttrs>();
   IndexExpr num_elem = tvm::cast(tvm::Int(32), tvm::ceil(
       tvm::cast(tvm::Float(32), param->stop - param->start) / param->step));
-  num_elem = tvm::ir::Simplify(num_elem);
   reporter->Assign(types[0], TensorTypeNode::make({num_elem}, param->dtype));
   return true;
 }

tests/python/relay/test_op_level3.py

@@ -460,7 +460,6 @@ def test_infer_type_prelu():
 
 def test_arange():
     def verify_arange(start, stop, step):
-        print(start, stop, step)
         dtype = "float32"
         if start is None and step is None:
             x = relay.arange(stop)
@@ -494,26 +493,26 @@ def verify_arange(start, stop, step):
 
 
 if __name__ == "__main__":
-    # test_cast()
-    # test_zeros_ones()
-    # test_unary_identity()
-    # test_clip()
-    # test_transpose_infer_type()
-    # test_transpose()
-    # test_reshape_infer_type()
-    # test_reshape()
-    # test_reshape_like_infer_type()
-    # test_reshape_like()
-    # test_take_infer_type()
-    # test_take()
-    # test_full_infer_type()
-    # test_full()
-    # test_full_like_infer_type()
-    # test_full_like()
-    # test_infer_type_leaky_relu()
-    # test_infer_type_prelu()
-    # test_squeeze()
-    # test_squeeze_infer_type()
-    # test_squeeze_bad_axes_infer_type()
-    # test_split_infer_type()
+    test_cast()
+    test_zeros_ones()
+    test_unary_identity()
+    test_clip()
+    test_transpose_infer_type()
+    test_transpose()
+    test_reshape_infer_type()
+    test_reshape()
+    test_reshape_like_infer_type()
+    test_reshape_like()
+    test_take_infer_type()
+    test_take()
+    test_full_infer_type()
+    test_full()
+    test_full_like_infer_type()
+    test_full_like()
+    test_infer_type_leaky_relu()
+    test_infer_type_prelu()
+    test_squeeze()
+    test_squeeze_infer_type()
+    test_squeeze_bad_axes_infer_type()
+    test_split_infer_type()
     test_arange()

topi/python/topi/transform.py

@@ -292,6 +292,28 @@ def tensordot(a, b, axes):
 
 
 def arange(stop, start=None, step=1, dtype="float32"):
+    """Creates a tensor with evenly spaced values within a given interval.
+
+    Parameters
+    ----------
+    stop : tvm.Expr
+        Stop of interval. The interval does not include this value.
+
+    start : tvm.Expr, optional
+        Start of interval. The interval includes this value. The default start
+        value is 0.
+
+    step : tvm.Expr, optional
+        Spacing between values. The default step size is 1.
+
+    dtype : str, optional
+        The target data type.
+
+    Returns
+    -------
+    result : tvm.Tensor
+        The resulting tensor.
+    """
     if start is None:
         start = 0
     else: