From 2eeaec7512a184bc510460659b22a73d216399d3 Mon Sep 17 00:00:00 2001
From: Bairong <zhuangbairong@huawei.com>
Date: Fri, 4 Dec 2020 00:40:21 +0800
Subject: [PATCH] support arange in functional.py

---
 mindspore/ops/functional.py          | 6 ++++++
 tests/ut/python/ops/test_math_ops.py | 9 +++++++++
 2 files changed, 15 insertions(+)

diff --git a/mindspore/ops/functional.py b/mindspore/ops/functional.py
index e9e3d4aa1d..00aaf64e49 100644
--- a/mindspore/ops/functional.py
+++ b/mindspore/ops/functional.py
@@ -22,6 +22,12 @@ from .primitive import Primitive
 from . import operations as P
 from .operations import _grad_ops
 
+
+def arange(start, limit=None, delta=1):
+    from mindspore.nn import Range
+    return Range(start, limit, delta)()
+
+
 typeof = Primitive('typeof')
 hastype = Primitive('hastype')
 cast = P.Cast()
diff --git a/tests/ut/python/ops/test_math_ops.py b/tests/ut/python/ops/test_math_ops.py
index c1c8408627..934b9f139c 100755
--- a/tests/ut/python/ops/test_math_ops.py
+++ b/tests/ut/python/ops/test_math_ops.py
@@ -24,6 +24,7 @@ from mindspore import Tensor
 from mindspore.common import dtype as mstype
 from mindspore.ops import composite as C
 from mindspore.ops import operations as P
+from mindspore.ops import functional as F
 from mindspore.ops import prim_attr_register, PrimitiveWithInfer
 from ..ut_filter import non_graph_engine
 from ....mindspore_test_framework.mindspore_test import mindspore_test
@@ -41,6 +42,7 @@ context.set_context(mode=context.GRAPH_MODE)
 
 grad = C.GradOperation()
 
+
 def test_multiply():
     """ test_multiply """
     input_x = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]))
@@ -138,6 +140,13 @@ def test_eye():
     assert np.all(eye_output.asnumpy() == expect)
 
 
+def test_arange():
+    """ test_arange """
+    F.arange(10)
+    F.arange(1, 5)
+    F.arange(1, 10, 2)
+
+
 class VirtualLossGrad(PrimitiveWithInfer):
     """ VirtualLossGrad definition """