[PIR Unittest] fix 4 pir uts test_elementwise_(add_op,div_op,nn_grad,sub_op) (#67403)

ooooo-create · web-flow · commit ff57d6f22161 · 2024-08-15T11:05:44.000+08:00
diff --git a/test/legacy_test/test_elementwise_add_op.py b/test/legacy_test/test_elementwise_add_op.py
@@ -620,12 +620,13 @@ def _executed_api(self, x, y, name=None):
         return paddle.add(x, y, name)
 
     def test_name(self):
-        with base.program_guard(base.Program()):
-            x = paddle.static.data(name="x", shape=[2, 3], dtype="float32")
-            y = paddle.static.data(name='y', shape=[2, 3], dtype='float32')
+        with paddle.pir_utils.OldIrGuard():
+            with base.program_guard(base.Program()):
+                x = paddle.static.data(name="x", shape=[2, 3], dtype="float32")
+                y = paddle.static.data(name='y', shape=[2, 3], dtype='float32')
 
-            y_1 = self._executed_api(x, y, name='add_res')
-            self.assertEqual(('add_res' in y_1.name), True)
+                y_1 = self._executed_api(x, y, name='add_res')
+                self.assertEqual(('add_res' in y_1.name), True)
 
     def test_declarative(self):
         with base.program_guard(base.Program()):
@@ -642,7 +643,7 @@ def gen_data():
 
             place = base.CPUPlace()
             exe = base.Executor(place)
-            z_value = exe.run(feed=gen_data(), fetch_list=[z.name])
+            z_value = exe.run(feed=gen_data(), fetch_list=[z])
             z_expected = np.array([3.0, 8.0, 6.0])
             self.assertEqual((z_value == z_expected).all(), True)
 
@@ -857,27 +858,30 @@ def test_float16_add(self):
 
 class TestTensorAddAPIWarnings(unittest.TestCase):
     def test_warnings(self):
-        with warnings.catch_warnings(record=True) as context:
-            warnings.simplefilter("always")
-
-            paddle.enable_static()
-            helper = LayerHelper("elementwise_add")
-            data = paddle.static.data(
-                name='data', shape=[None, 3, 32, 32], dtype='float32'
-            )
-            out = helper.create_variable_for_type_inference(dtype=data.dtype)
-            os.environ['FLAGS_print_extra_attrs'] = "1"
-            helper.append_op(
-                type="elementwise_add",
-                inputs={'X': data, 'Y': data},
-                outputs={'Out': out},
-                attrs={'axis': 1, 'use_mkldnn': False},
-            )
-            self.assertTrue(
-                "op elementwise_add's attr axis = 1 is not the default value: -1"
-                in str(context[-1].message)
-            )
-            os.environ['FLAGS_print_extra_attrs'] = "0"
+        with paddle.pir_utils.OldIrGuard():
+            with warnings.catch_warnings(record=True) as context:
+                warnings.simplefilter("always")
+
+                paddle.enable_static()
+                helper = LayerHelper("elementwise_add")
+                data = paddle.static.data(
+                    name='data', shape=[None, 3, 32, 32], dtype='float32'
+                )
+                out = helper.create_variable_for_type_inference(
+                    dtype=data.dtype
+                )
+                os.environ['FLAGS_print_extra_attrs'] = "1"
+                helper.append_op(
+                    type="elementwise_add",
+                    inputs={'X': data, 'Y': data},
+                    outputs={'Out': out},
+                    attrs={'axis': 1, 'use_mkldnn': False},
+                )
+                self.assertTrue(
+                    "op elementwise_add's attr axis = 1 is not the default value: -1"
+                    in str(context[-1].message)
+                )
+                os.environ['FLAGS_print_extra_attrs'] = "0"
 
 
 class TestTensorFloat32Bfloat16OrFloat16Add(unittest.TestCase):
diff --git a/test/legacy_test/test_elementwise_div_op.py b/test/legacy_test/test_elementwise_div_op.py
@@ -18,6 +18,7 @@
 from op_test import OpTest, convert_float_to_uint16, skip_check_grad_ci
 
 import paddle
+import paddle.static
 from paddle import base
 from paddle.base import core
 from paddle.pir_utils import test_with_pir_api
@@ -509,14 +510,15 @@ def test_shape_with_batch_sizes(self):
 class TestDivideOp(unittest.TestCase):
     def test_name(self):
         paddle.enable_static()
-        main_program = paddle.static.Program()
-        with paddle.static.program_guard(main_program):
-            x = paddle.static.data(name="x", shape=[2, 3], dtype="float32")
-            y = paddle.static.data(name='y', shape=[2, 3], dtype='float32')
+        with paddle.pir_utils.OldIrGuard():
+            main_program = paddle.static.Program()
+            with paddle.static.program_guard(main_program):
+                x = paddle.static.data(name="x", shape=[2, 3], dtype="float32")
+                y = paddle.static.data(name='y', shape=[2, 3], dtype='float32')
 
-            y_1 = paddle.divide(x, y, name='div_res')
+                y_1 = paddle.divide(x, y, name='div_res')
 
-            self.assertEqual(('div_res' in y_1.name), True)
+                self.assertEqual(('div_res' in y_1.name), True)
 
         paddle.disable_static()
 
diff --git a/test/legacy_test/test_elementwise_nn_grad.py b/test/legacy_test/test_elementwise_nn_grad.py
@@ -552,7 +552,8 @@ def test_grad(self):
         if core.is_compiled_with_cuda():
             places.append(base.CUDAPlace(0))
         for p in places:
-            self.func(p)
+            with paddle.pir_utils.OldIrGuard():
+                self.func(p)
 
 
 class TestElementwiseMulTripleGradCheck(unittest.TestCase):
@@ -621,7 +622,8 @@ def test_grad(self):
         if core.is_compiled_with_cuda():
             places.append(base.CUDAPlace(0))
         for p in places:
-            self.func(p)
+            with paddle.pir_utils.OldIrGuard():
+                self.func(p)
 
 
 if __name__ == "__main__":
diff --git a/test/legacy_test/test_elementwise_sub_op.py b/test/legacy_test/test_elementwise_sub_op.py
@@ -897,12 +897,13 @@ def _executed_api(self, x, y, name=None):
         return paddle.subtract(x, y, name)
 
     def test_name(self):
-        with base.program_guard(base.Program()):
-            x = paddle.static.data(name="x", shape=[2, 3], dtype="float32")
-            y = paddle.static.data(name='y', shape=[2, 3], dtype=np.float32)
+        with paddle.pir_utils.OldIrGuard():
+            with base.program_guard(base.Program()):
+                x = paddle.static.data(name="x", shape=[2, 3], dtype="float32")
+                y = paddle.static.data(name='y', shape=[2, 3], dtype=np.float32)
 
-            y_1 = self._executed_api(x, y, name='subtract_res')
-            self.assertEqual(('subtract_res' in y_1.name), True)
+                y_1 = self._executed_api(x, y, name='subtract_res')
+                self.assertEqual(('subtract_res' in y_1.name), True)
 
     @test_with_pir_api
     def test_declarative(self):
@@ -1063,27 +1064,30 @@ def test_dygraph_sub(self):
 
 class TestTensorSubAPIWarnings(unittest.TestCase):
     def test_warnings(self):
-        with warnings.catch_warnings(record=True) as context:
-            warnings.simplefilter("always")
-
-            paddle.enable_static()
-            helper = LayerHelper("elementwise_sub")
-            data = paddle.static.data(
-                name='data', shape=[None, 3, 32, 32], dtype=np.float32
-            )
-            out = helper.create_variable_for_type_inference(dtype=data.dtype)
-            os.environ['FLAGS_print_extra_attrs'] = "1"
-            helper.append_op(
-                type="elementwise_sub",
-                inputs={'X': data, 'Y': data},
-                outputs={'Out': out},
-                attrs={'axis': 1, 'use_mkldnn': False},
-            )
-            self.assertTrue(
-                "op elementwise_sub's attr axis = 1 is not the default value: -1"
-                in str(context[-1].message)
-            )
-            os.environ['FLAGS_print_extra_attrs'] = "0"
+        with paddle.pir_utils.OldIrGuard():
+            with warnings.catch_warnings(record=True) as context:
+                warnings.simplefilter("always")
+
+                paddle.enable_static()
+                helper = LayerHelper("elementwise_sub")
+                data = paddle.static.data(
+                    name='data', shape=[None, 3, 32, 32], dtype=np.float32
+                )
+                out = helper.create_variable_for_type_inference(
+                    dtype=data.dtype
+                )
+                os.environ['FLAGS_print_extra_attrs'] = "1"
+                helper.append_op(
+                    type="elementwise_sub",
+                    inputs={'X': data, 'Y': data},
+                    outputs={'Out': out},
+                    attrs={'axis': 1, 'use_mkldnn': False},
+                )
+                self.assertTrue(
+                    "op elementwise_sub's attr axis = 1 is not the default value: -1"
+                    in str(context[-1].message)
+                )
+                os.environ['FLAGS_print_extra_attrs'] = "0"
 
 
 if __name__ == '__main__':
