[PIR] reopen TestReduceAPI.test_static_reduce UT (#63185)

gouzil · SigureMo · web-flow · commit 0891c1838e5b · 2024-04-07T11:14:22.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -389,7 +389,6 @@ def remove_op(block, op, state):
     '''
     remove op from block
     '''
-    block.remove_op(op)
     if state.opgrad_to_op[op] != []:
         fwd_op = state.opgrad_to_op[op][0]
         state.op_to_opgrad[fwd_op].remove(op)
@@ -403,6 +402,10 @@ def remove_op(block, op, state):
                 raise ValueError(
                     'input_grad in [%s] is value which need to sum ', op.name()
                 )
+    # NOTE(SigureMo): Ensure access to the op's results before removing it.
+    # Otherwise, the op will be deconstructed and access the num_results
+    # will be undefined behavior, it always cause hanging on the macOS.
+    block.remove_op(op)
 
 
 def while_prune_check(while_tuple_ops):
diff --git a/test/legacy_test/test_zero_dim_reduce_api.py b/test/legacy_test/test_zero_dim_reduce_api.py
@@ -22,6 +22,7 @@
 import numpy as np
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 reduce_api_list = [
     paddle.sum,
@@ -142,13 +143,11 @@ def test_dygraph_reduce(self):
 
         paddle.enable_static()
 
-    # TODO(SigureMo): Temporarily disable this test case in due to hanging in mac CI.
-    # @test_with_pir_api
-    def test_static_reduce(self):
+    @test_with_pir_api
+    def test_static_reduce_x_0D(self):
         paddle.enable_static()
         for api in reduce_api_list:
             main_prog = paddle.static.Program()
-            block = main_prog.global_block()
             exe = paddle.static.Executor()
             with paddle.static.program_guard(
                 main_prog, paddle.static.Program()
@@ -188,17 +187,26 @@ def test_static_reduce(self):
                 )
                 res = exe.run(main_prog, fetch_list=fetch_list)
 
-                self.assertEqual(res[0].shape, ())
-                self.assertEqual(res[1].shape, ())
+                for res_data in res:
+                    self.assertEqual(res_data.shape, ())
                 if api not in [paddle.count_nonzero]:
                     np.testing.assert_allclose(res[0], res[1])
 
+                if len(res) > 3:
+                    np.testing.assert_allclose(res[-2], np.array(1.0))
+                    np.testing.assert_allclose(res[-1], np.array(1.0))
                 if len(res) > 2:
-                    self.assertEqual(res[2].shape, ())
-                    self.assertEqual(res[3].shape, ())
-                    np.testing.assert_allclose(res[2], np.array(1.0))
-                    np.testing.assert_allclose(res[3], np.array(1.0))
+                    np.testing.assert_allclose(res[-1], np.array(1.0))
 
+    @test_with_pir_api
+    def test_static_reduce_ND_0D(self):
+        paddle.enable_static()
+        for api in reduce_api_list:
+            main_prog = paddle.static.Program()
+            exe = paddle.static.Executor()
+            with paddle.static.program_guard(
+                main_prog, paddle.static.Program()
+            ):
                 # 2) x is ND, reduce to 0D
                 if api in [paddle.all, paddle.any]:
                     x = paddle.randint(0, 2, [3, 5]).astype('bool')
@@ -229,6 +237,15 @@ def test_static_reduce(self):
                 if len(res) > 2:
                     self.assertEqual(res[2].shape, (3, 5))
 
+    @test_with_pir_api
+    def test_static_reduce_x_1D(self):
+        paddle.enable_static()
+        for api in reduce_api_list:
+            main_prog = paddle.static.Program()
+            exe = paddle.static.Executor()
+            with paddle.static.program_guard(
+                main_prog, paddle.static.Program()
+            ):
                 # 3) x is 1D, axis=0, reduce to 0D
                 if api in [paddle.all, paddle.any]:
                     x = paddle.randint(0, 2, [5]).astype('bool')
