[PIR] Migrate uniform to pir (#57433)

changeyoung98 · web-flow · commit 3e7173687391 · 2023-09-19T10:22:15.000+08:00
* uniform new_ir op_test passed 8/9

* close new_ir check for tanh complex
diff --git a/python/paddle/tensor/random.py b/python/paddle/tensor/random.py
@@ -18,7 +18,7 @@
 from paddle import _C_ops, _legacy_C_ops
 from paddle.base.framework import _current_expected_place
 from paddle.common_ops_import import Variable
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_mode, in_dynamic_or_pir_mode
 
 from ..base.data_feeder import (
     check_dtype,
@@ -645,7 +645,7 @@ def uniform(shape, dtype=None, min=-1.0, max=1.0, seed=0, name=None):
     if not isinstance(dtype, core.VarDesc.VarType):
         dtype = convert_np_dtype_to_dtype_(dtype)
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         shape = paddle.utils.convert_shape_to_list(shape)
         return _C_ops.uniform(
             shape,
diff --git a/python/paddle/utils/layers_utils.py b/python/paddle/utils/layers_utils.py
@@ -455,7 +455,10 @@ def convert_shape_to_list(shape):
     if isinstance(shape, (list, tuple)):
         shape = [x.item(0) if isinstance(x, Variable) else x for x in shape]
     else:
-        shape = shape.astype(int).tolist()
+        if in_dygraph_mode():
+            shape = shape.astype(int).tolist()
+        else:
+            shape = [shape]
     return shape
 
 
diff --git a/test/legacy_test/op_test.py b/test/legacy_test/op_test.py
@@ -1624,9 +1624,7 @@ def _construct_grad_program_from_forward(
             if fwd_var_name is None:
                 fwd_var_name = arg
             fwd_var = fwd_program.global_block().vars.get(fwd_var_name)
-            assert fwd_var is not None, "{} cannot be found".format(
-                fwd_var_name
-            )
+            assert fwd_var is not None, f"{fwd_var_name} cannot be found"
             grad_var = grad_block.create_var(
                 name=arg,
                 dtype=fwd_var.dtype,
@@ -2682,7 +2680,9 @@ def check_output(
                     return
                 self.check_compile_vs_runtime(fetch_list, outs)
 
-    def check_output_customized(self, checker, custom_place=None):
+    def check_output_customized(
+        self, checker, custom_place=None, check_new_ir=False
+    ):
         self.__class__.op_type = self.op_type
         places = self._get_places()
         if custom_place:
@@ -2692,6 +2692,12 @@ def check_output_customized(self, checker, custom_place=None):
             outs = [np.array(out) for out in outs]
             outs.sort(key=len)
             checker(outs)
+            if check_new_ir:
+                with paddle.new_ir_utils.IrGuard():
+                    outs_p = self._calc_new_ir_output(place)
+                    outs_p = [outs_p[out] for out in outs_p]
+                    outs_p.sort(key=len)
+                    checker(outs_p)
 
     def check_output_with_place_customized(self, checker, place):
         outs = self.calc_output(place)
diff --git a/test/legacy_test/test_activation_op.py b/test/legacy_test/test_activation_op.py
@@ -692,7 +692,7 @@ def test_check_grad(self):
             return
         # TODO(ScottWong98): set `check_prim=False` when `fill_any_like` supports `complex` dtype
         if self.dtype == np.complex64 or self.dtype == np.complex128:
-            self.check_grad(['X'], 'Out', check_prim=False, check_new_ir=True)
+            self.check_grad(['X'], 'Out', check_prim=False, check_new_ir=False)
         else:
             self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
diff --git a/test/legacy_test/test_uniform_random_op.py b/test/legacy_test/test_uniform_random_op.py
@@ -69,7 +69,7 @@ def init_attrs(self):
         self.output_hist = output_hist
 
     def test_check_output(self):
-        self.check_output_customized(self.verify_output)
+        self.check_output_customized(self.verify_output, check_new_ir=True)
 
     def verify_output(self, outs):
         hist, prob = self.output_hist(np.array(outs[0]))
@@ -101,7 +101,7 @@ def init_attrs(self):
         self.output_hist = output_hist
 
     def test_check_output(self):
-        self.check_output_customized(self.verify_output)
+        self.check_output_customized(self.verify_output, check_new_ir=True)
 
     def verify_output(self, outs):
         hist, prob = self.output_hist(np.array(outs[0]))
@@ -121,7 +121,7 @@ def init_attrs(self):
         self.output_hist = output_hist
 
     def test_check_output(self):
-        self.check_output_customized(self.verify_output)
+        self.check_output_customized(self.verify_output, check_new_ir=True)
 
     def verify_output(self, outs):
         hist, prob = self.output_hist(np.array(outs[0]))
@@ -141,7 +141,7 @@ def init_attrs(self):
         self.output_hist = output_hist
 
     def test_check_output(self):
-        self.check_output_customized(self.verify_output)
+        self.check_output_customized(self.verify_output, check_new_ir=True)
 
     def verify_output(self, outs):
         hist, prob = self.output_hist(np.array(outs[0]))
@@ -170,7 +170,7 @@ def init_attrs(self):
         self.output_hist = output_hist
 
     def test_check_output(self):
-        self.check_output_customized(self.verify_output)
+        self.check_output_customized(self.verify_output, check_new_ir=True)
 
     def verify_output(self, outs):
         hist, prob = self.output_hist(np.array(outs[0]))
@@ -203,6 +203,7 @@ def init_dtype(self):
 
 class TestUniformRandomOpError(unittest.TestCase):
     def test_errors(self):
+        paddle.enable_static()
         main_prog = Program()
         start_prog = Program()
         with program_guard(main_prog, start_prog):
@@ -226,6 +227,7 @@ def test_out_dtype():
                 self.assertEqual(out.dtype, base.core.VarDesc.VarType.FP64)
 
             test_out_dtype()
+        paddle.disable_static()
 
 
 class TestUniformRandomOpWithDiagInit(TestUniformRandomOp):
@@ -241,6 +243,9 @@ def init_attrs(self):
         }
         self.output_hist = output_hist_diag
 
+    def test_check_output(self):
+        self.check_output_customized(self.verify_output, check_new_ir=False)
+
 
 class TestUniformRandomOpSelectedRows(unittest.TestCase):
     def get_places(self):
@@ -297,6 +302,7 @@ def check_with_place(self, place):
 
 class TestUniformRandomOpApi(unittest.TestCase):
     def test_api(self):
+        paddle.enable_static()
         paddle.seed(10)
         x = paddle.static.data(
             'x', shape=[-1, 16], dtype='float32', lod_level=1
@@ -321,10 +327,12 @@ def test_api(self):
         exe = base.Executor(place)
         exe.run(base.default_startup_program())
         ret = exe.run(feed={'x': x_tensor}, fetch_list=[y], return_numpy=False)
+        paddle.disable_static()
 
 
 class TestUniformRandomOp_attr_tensor_API(unittest.TestCase):
     def test_attr_tensor_API(self):
+        paddle.enable_static()
         startup_program = base.Program()
         train_program = base.Program()
         with base.program_guard(train_program, startup_program):
@@ -338,8 +346,10 @@ def test_attr_tensor_API(self):
 
             exe.run(startup_program)
             outs = exe.run(train_program, fetch_list=[ret])
+        paddle.disable_static()
 
     def test_attr_tensorlist_int32_API(self):
+        paddle.enable_static()
         startup_program = base.Program()
         train_program = base.Program()
         with base.program_guard(train_program, startup_program):
@@ -354,8 +364,10 @@ def test_attr_tensorlist_int32_API(self):
 
             exe.run(startup_program)
             outs = exe.run(train_program, fetch_list=[ret])
+        paddle.disable_static()
 
     def test_attr_tensor_int32_API(self):
+        paddle.enable_static()
         startup_program = base.Program()
         train_program = base.Program()
         with base.program_guard(train_program, startup_program):
@@ -373,10 +385,12 @@ def test_attr_tensor_int32_API(self):
             outs = exe.run(
                 train_program, feed={'shape_tensor': Shape}, fetch_list=[ret]
             )
+        paddle.disable_static()
 
 
 class TestUniformRandomOp_API_seed(unittest.TestCase):
     def test_attr_tensor_API(self):
+        paddle.enable_static()
         _seed = 10
         gen = paddle.seed(_seed)
         startup_program = base.Program()
@@ -399,6 +413,7 @@ def test_attr_tensor_API(self):
             for i in ret_value.flatten():
                 self.assertGreaterEqual(i, _min)
                 self.assertLess(i, _max)
+        paddle.disable_static()
 
 
 class TestUniformRandomOpSelectedRowsShapeTensor(unittest.TestCase):
@@ -476,6 +491,7 @@ def test_check_output(self):
 
 class TestUniformRandomBatchSizeLikeOpError(unittest.TestCase):
     def test_errors(self):
+        paddle.enable_static()
         main_prog = Program()
         start_prog = Program()
         with program_guard(main_prog, start_prog):
@@ -503,6 +519,7 @@ def test_dtype():
                 random.uniform_random_batch_size_like(x2, 'int32')
 
             self.assertRaises(TypeError, test_dtype)
+        paddle.disable_static()
 
 
 class TestUniformAlias(unittest.TestCase):
@@ -519,6 +536,7 @@ def test_uniform_random():
 
 class TestUniformOpError(unittest.TestCase):
     def test_errors(self):
+        paddle.enable_static()
         main_prog = Program()
         start_prog = Program()
         with program_guard(main_prog, start_prog):
@@ -552,6 +570,7 @@ def test_out_dtype():
                 self.assertEqual(out.dtype, base.core.VarDesc.VarType.FP64)
 
             test_out_dtype()
+        paddle.disable_static()
 
 
 class TestUniformDygraphMode(unittest.TestCase):
