[PIR] No.31 Migrate paddle.arange into pir (#57344)

BeingGod · web-flow · commit 7a9669bb2fa7 · 2023-09-22T23:15:25.000+08:00
diff --git a/python/paddle/tensor/attribute.py b/python/paddle/tensor/attribute.py
@@ -20,7 +20,11 @@
 from paddle import _C_ops
 
 from ..base.data_feeder import check_type, check_variable_and_dtype
-from ..base.framework import in_dygraph_mode, in_dynamic_or_pir_mode
+from ..base.framework import (
+    in_dygraph_mode,
+    in_dynamic_or_pir_mode,
+    in_pir_mode,
+)
 from ..common_ops_import import Variable
 from ..framework import LayerHelper, core
 from .creation import _complex_to_real_dtype, assign
@@ -233,16 +237,30 @@ def is_integer(x):
             >>> print(paddle.is_integer(x))
             True
     """
-    if not isinstance(x, (paddle.Tensor, paddle.static.Variable)):
+    if not isinstance(
+        x, (paddle.Tensor, paddle.static.Variable, paddle.pir.OpResult)
+    ):
         raise TypeError(f"Expected Tensor, but received type of x: {type(x)}")
     dtype = x.dtype
-    is_int_dtype = (
-        dtype == core.VarDesc.VarType.UINT8
-        or dtype == core.VarDesc.VarType.INT8
-        or dtype == core.VarDesc.VarType.INT16
-        or dtype == core.VarDesc.VarType.INT32
-        or dtype == core.VarDesc.VarType.INT64
-    )
+
+    is_int_dtype = False
+    if not in_pir_mode():
+        is_int_dtype = (
+            dtype == core.VarDesc.VarType.UINT8
+            or dtype == core.VarDesc.VarType.INT8
+            or dtype == core.VarDesc.VarType.INT16
+            or dtype == core.VarDesc.VarType.INT32
+            or dtype == core.VarDesc.VarType.INT64
+        )
+    else:
+        is_int_dtype = (
+            dtype == core.DataType.INT8
+            or dtype == core.DataType.INT8
+            or dtype == core.DataType.INT16
+            or dtype == core.DataType.INT32
+            or dtype == core.DataType.INT64
+        )
+
     return is_int_dtype
 
 
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -1343,8 +1343,8 @@ def arange(start=0, end=None, step=1, dtype=None, name=None):
 
     if dtype is None:
         for val in [start, end, step]:
-            if isinstance(val, Variable):
-                if not val.is_integer():
+            if isinstance(val, (Variable, paddle.pir.OpResult)):
+                if not paddle.is_integer(val):
                     dtype = paddle.get_default_dtype()
                     break
                 else:
@@ -1357,35 +1357,35 @@ def arange(start=0, end=None, step=1, dtype=None, name=None):
                     dtype = 'int64'
 
     out_shape = None
-    if not in_dynamic_mode() and (
-        not isinstance(start, Variable)
-        and not isinstance(end, Variable)
-        and not isinstance(step, Variable)
+    if not in_dynamic_or_pir_mode() and (
+        not isinstance(start, (Variable, paddle.pir.OpResult))
+        and not isinstance(end, (Variable, paddle.pir.OpResult))
+        and not isinstance(step, (Variable, paddle.pir.OpResult))
     ):
         out_shape = [int(math.ceil((end - start) / step))]
 
-    if not isinstance(dtype, core.VarDesc.VarType):
+    if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
         dtype = convert_np_dtype_to_dtype_(dtype)
 
-    if not isinstance(start, Variable):
+    if not isinstance(start, (Variable, paddle.pir.OpResult)):
         with device_guard("cpu"):
             start = fill_constant([1], dtype, start, force_cpu=True)
     elif start.dtype != dtype:
         start = paddle.cast(start, dtype)
 
-    if not isinstance(end, Variable):
+    if not isinstance(end, (Variable, paddle.pir.OpResult)):
         with device_guard("cpu"):
             end = fill_constant([1], dtype, end, force_cpu=True)
     elif end.dtype != dtype:
         end = paddle.cast(end, dtype)
 
-    if not isinstance(step, Variable):
+    if not isinstance(step, (Variable, paddle.pir.OpResult)):
         with device_guard("cpu"):
             step = fill_constant([1], dtype, step, force_cpu=True)
     elif step.dtype != dtype:
         step = paddle.cast(step, dtype)
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.arange(start, end, step, dtype, _current_expected_place())
     else:
         check_dtype(
diff --git a/test/legacy_test/test_arange.py b/test/legacy_test/test_arange.py
@@ -48,7 +48,7 @@ def init_config(self):
         self.case = (0, 1, 0.2)
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 class TestFloatArangeOp(TestArangeOp):
@@ -65,7 +65,7 @@ def init_config(self):
         self.case = (0, 5, 1)
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
 
 @unittest.skipIf(
@@ -99,7 +99,7 @@ def init_config(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place)
+        self.check_output_with_place(place, check_new_ir=True)
 
 
 class TestInt32ArangeOp(TestArangeOp):
@@ -131,13 +131,15 @@ def init_config(self):
 
 
 class TestArangeOpError(unittest.TestCase):
-    def test_errors(self):
+    def test_static_errors(self):
         with program_guard(Program(), Program()):
+            paddle.enable_static()
             self.assertRaises(TypeError, paddle.arange, 10, dtype='int8')
 
 
 class TestArangeAPI(unittest.TestCase):
     def test_out(self):
+        paddle.enable_static()
         with program_guard(Program(), Program()):
             x1 = paddle.arange(0, 5, 1, 'float32')
 
@@ -152,6 +154,7 @@ def test_out(self):
         expected_data = np.arange(0, 5, 1).astype(np.float32)
         self.assertEqual((out == expected_data).all(), True)
         self.assertListEqual(list(x1.shape), [5])
+        paddle.disable_static(place)
 
 
 class TestArangeImperative(unittest.TestCase):
