[PIR] support normal and fix TestNoBackwardAPIStatic.test_normal UT (#62864)

gouzil · web-flow · commit b809787eb743 · 2024-03-21T14:42:08.000+08:00
diff --git a/python/paddle/tensor/random.py b/python/paddle/tensor/random.py
@@ -741,18 +741,22 @@ def normal(mean=0.0, std=1.0, shape=None, name=None):
             [0.48646951, 0.00815189, 3.74022293])
             >>> # doctest: -SKIP
     """
-    if not in_dynamic_or_pir_mode():
-        check_type(mean, 'mean', (int, float, Variable), 'normal')
-        check_type(std, 'std', (int, float, Variable), 'normal')
-        if isinstance(mean, Variable):
+    if not in_dynamic_mode():
+        check_type(
+            mean, 'mean', (int, float, Variable, paddle.pir.Value), 'normal'
+        )
+        check_type(
+            std, 'std', (int, float, Variable, paddle.pir.Value), 'normal'
+        )
+        if isinstance(mean, (Variable, paddle.pir.Value)):
             check_dtype(
                 mean.dtype,
                 'mean',
                 ['float32', 'float64'],
                 'normal',
                 "If mean is Tensor, it's data type only support float32, float64.",
             )
-        if isinstance(std, Variable):
+        if isinstance(std, (Variable, paddle.pir.Value)):
             check_dtype(
                 std.dtype,
                 'std',
@@ -763,16 +767,16 @@ def normal(mean=0.0, std=1.0, shape=None, name=None):
         if shape is not None:
             check_shape(shape, 'normal')
 
-    if isinstance(mean, Variable):
-        if isinstance(std, Variable):
+    if isinstance(mean, (Variable, paddle.pir.Value)):
+        if isinstance(std, (Variable, paddle.pir.Value)):
             if std.dtype != mean.dtype:
                 std = paddle.cast(std, mean.dtype)
             mean_shape = paddle.shape(mean)
             std = paddle.reshape(std, mean_shape)
         else:
             std = float(std)
         out = standard_normal(paddle.shape(mean), mean.dtype, name)
-    elif isinstance(std, Variable):
+    elif isinstance(std, (Variable, paddle.pir.Value)):
         mean = float(mean)
         out = standard_normal(paddle.shape(std), std.dtype, name)
     else:
diff --git a/test/legacy_test/test_normal.py b/test/legacy_test/test_normal.py
@@ -18,6 +18,7 @@
 import numpy as np
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 np.random.seed(10)
 paddle.seed(10)
@@ -62,10 +63,11 @@ def static_api(self):
         ret_all_shape = copy.deepcopy(shape)
         ret_all_shape.insert(0, self.repeat_num)
         ret_all = np.zeros(ret_all_shape, self.dtype)
+        main_program = paddle.static.Program()
         if isinstance(self.mean, np.ndarray) and isinstance(
             self.std, np.ndarray
         ):
-            with paddle.static.program_guard(paddle.static.Program()):
+            with paddle.static.program_guard(main_program):
                 mean = paddle.static.data(
                     'Mean', self.mean.shape, self.mean.dtype
                 )
@@ -84,7 +86,7 @@ def static_api(self):
                     ret_all[i] = ret[0]
             return ret_all
         elif isinstance(self.mean, np.ndarray):
-            with paddle.static.program_guard(paddle.static.Program()):
+            with paddle.static.program_guard(main_program):
                 mean = paddle.static.data(
                     'Mean', self.mean.shape, self.mean.dtype
                 )
@@ -96,7 +98,7 @@ def static_api(self):
                     ret_all[i] = ret[0]
             return ret_all
         elif isinstance(self.std, np.ndarray):
-            with paddle.static.program_guard(paddle.static.Program()):
+            with paddle.static.program_guard(main_program):
                 std = paddle.static.data('Std', self.std.shape, self.std.dtype)
                 out = paddle.normal(self.mean, std, self.shape)
 
@@ -106,7 +108,7 @@ def static_api(self):
                     ret_all[i] = ret[0]
             return ret_all
         else:
-            with paddle.static.program_guard(paddle.static.Program()):
+            with paddle.static.program_guard(main_program):
                 out = paddle.normal(self.mean, self.std, self.shape)
 
                 exe = paddle.static.Executor(self.place)
@@ -138,6 +140,7 @@ def dygraph_api(self):
         paddle.enable_static()
         return ret_all
 
+    @test_with_pir_api
     def test_api(self):
         ret_static = self.static_api()
         ret_dygraph = self.dygraph_api()
@@ -185,6 +188,7 @@ def set_attrs(self):
 
 
 class TestNormalAlias(unittest.TestCase):
+    @test_with_pir_api
     def test_alias(self):
         paddle.disable_static()
         shape = [1, 2, 3]
@@ -195,8 +199,10 @@ def test_alias(self):
 
 
 class TestNormalErrors(unittest.TestCase):
+    @test_with_pir_api
     def test_errors(self):
-        with paddle.static.program_guard(paddle.static.Program()):
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
             mean = [1, 2, 3]
             self.assertRaises(TypeError, paddle.normal, mean)
 
diff --git a/test/legacy_test/test_zero_dim_no_backward_api.py b/test/legacy_test/test_zero_dim_no_backward_api.py
@@ -313,6 +313,7 @@ def test_arange(self):
         )[0]
         np.testing.assert_array_equal(res, [1.0, 2.0, 3.0, 4.0, 5.0])
 
+    @test_with_pir_api
     def test_normal(self):
         mean = paddle.full([], 0.0)
         std = paddle.full([], 0.0)
