[PIR] Open pir tests for cond/if op (#61455)

chen2016013 · web-flow · commit 50082948176a · 2024-02-04T11:07:24.000+08:00
* open test

* Update test_optimizer_grad.py

* Update test_return.py
diff --git a/test/dygraph_to_static/test_return.py b/test/dygraph_to_static/test_return.py
@@ -19,14 +19,14 @@
     Dy2StTestBase,
     enable_to_static_guard,
     test_ast_only,
+    test_legacy_and_pt_and_pir,
     test_legacy_only,
 )
 from ifelse_simple_func import dyfunc_with_if_else
 
 import paddle
 from paddle.base import core
 from paddle.jit.dy2static.utils import Dygraph2StaticException
-from paddle.pir_utils import test_with_pir_api
 
 SEED = 2020
 np.random.seed(SEED)
@@ -279,7 +279,7 @@ def _test_value_impl(self):
             self.assertEqual(dygraph_res, static_res)
 
     @test_ast_only
-    @test_with_pir_api
+    @test_legacy_and_pt_and_pir
     def test_transformed_static_result(self):
         self.init_dygraph_func()
         if hasattr(self, "error"):
@@ -326,8 +326,7 @@ def _test_value_impl(self):
         else:
             self.assertEqual(dygraph_res, static_res)
 
-    # Why add test_legacy_only? : PIR not support if true and false branch output with different dtype
-    @test_legacy_only
+    @test_legacy_and_pt_and_pir
     @test_ast_only
     def test_transformed_static_result(self):
         self.init_dygraph_func()
@@ -429,8 +428,7 @@ def _test_value_impl(self):
         else:
             self.assertEqual(dygraph_res, static_res)
 
-    # Why add test_legacy_only? : PIR not support if true and false branch output with different dtype
-    @test_legacy_only
+    @test_legacy_and_pt_and_pir
     @test_ast_only
     def test_transformed_static_result(self):
         self.init_dygraph_func()
