PaddlePaddle
diff --git a/‎python/paddle/jit/dy2static/convert_operators.py
Copy file name to clipboard
+5-1Lines changed: 5 additions & 1 deletion b/‎python/paddle/jit/dy2static/convert_operators.py
Copy file name to clipboard
+5-1Lines changed: 5 additions & 1 deletion
diff --git a/‎python/paddle/jit/pir_dy2static/parameter_recorder.py
Copy file name to clipboard
+5-1Lines changed: 5 additions & 1 deletion b/‎python/paddle/jit/pir_dy2static/parameter_recorder.py
Copy file name to clipboard
+5-1Lines changed: 5 additions & 1 deletion
diff --git a/‎python/paddle/pir/math_op_patch.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎python/paddle/pir/math_op_patch.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎test/dygraph_to_static/check_approval.py
Copy file name to clipboard
+1-31Lines changed: 1 addition & 31 deletions b/‎test/dygraph_to_static/check_approval.py
Copy file name to clipboard
+1-31Lines changed: 1 addition & 31 deletions
diff --git a/‎test/dygraph_to_static/dygraph_to_static_utils.py
Copy file name to clipboard
+11-4Lines changed: 11 additions & 4 deletions b/‎test/dygraph_to_static/dygraph_to_static_utils.py
Copy file name to clipboard
+11-4Lines changed: 11 additions & 4 deletions
diff --git a/‎test/dygraph_to_static/test_cpu_cuda_to_tensor.py
Copy file name to clipboard
+1-39Lines changed: 1 addition & 39 deletions b/‎test/dygraph_to_static/test_cpu_cuda_to_tensor.py
Copy file name to clipboard
+1-39Lines changed: 1 addition & 39 deletions
diff --git a/‎test/dygraph_to_static/test_incubate_jit_inference.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎test/dygraph_to_static/test_incubate_jit_inference.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎test/dygraph_to_static/test_op_attr.py
Copy file name to clipboard
+8-1Lines changed: 8 additions & 1 deletion b/‎test/dygraph_to_static/test_op_attr.py
Copy file name to clipboard
+8-1Lines changed: 8 additions & 1 deletion
diff --git a/‎test/dygraph_to_static/test_partial_program_hook.py
Copy file name to clipboard
-58Lines changed: 0 additions & 58 deletions b/‎test/dygraph_to_static/test_partial_program_hook.py
Copy file name to clipboard
-58Lines changed: 0 additions & 58 deletions
@@ -447,8 +447,12 @@ def _run_paddle_cond(
     pred = cast_bool_if_necessary(pred)
     init_args = helper.get(return_name_ids)
     from paddle.jit.dy2static.program_translator import ProgramTranslator
+    from paddle.jit.pir_dy2static.parameter_recorder import _global_inplace_map
 
-    inplace_map = ProgramTranslator.get_instance()._inplace_map
+    if use_pir_api():
+        inplace_map = _global_inplace_map
+    else:
+        inplace_map = ProgramTranslator.get_instance()._inplace_map
     union_name = None
     # TODO(@xiongkun) lambda can have push_pop_names, which will cause error.
     if return_name_ids is None and push_pop_names is None:
 
@@ -110,7 +110,11 @@ def restore_checkpoint(self, checkpoint):
         self.params_dict = checkpoint
 
     def save_checkpoint(self):
-        return dict(self.params_dict.items())
+        ckeckpoint = {}
+        for program_id, params in self.params_dict.items():
+            new_params = dict(params.items())
+            ckeckpoint[program_id] = new_params
+        return ckeckpoint
 
 
 _global_parameter_recorder = ParametersRecorder()
 
@@ -912,7 +912,7 @@ def numpy(self):
         pass
 
     @fake_interface_only
-    def register_hook(self):
+    def register_hook(self, hook):
         """
         Value don't have 'register_hook' interface in static graph mode
         But this interface can greatly facilitate dy2static.
 
@@ -105,21 +105,12 @@ def __init__(self, start: Location, end: Location):
         )
 
 
-class TestCaseWithoutDecoratorDiagnostic(Diagnostic):
-    def __init__(self, start: Location, end: Location):
-        super().__init__(
-            start,
-            end,
-            'Test case should use @test_legacy_and_pt_and_pir instead of no decorator',
-        )
-
-
 class TestCaseWithPirApiDecoratorDiagnostic(Diagnostic):
     def __init__(self, start: Location, end: Location):
         super().__init__(
             start,
             end,
-            'Test case should use @test_legacy_and_pt_and_pir instead of @test_with_pir_api',
+            'Test case should use @test_pt_and_pir instead of @test_with_pir_api',
         )
 
 
@@ -151,19 +142,6 @@ def __init__(self, start: Location, end: Location):
         "test_pylayer.py",
         "test_tensor_attr_consistency.py",
     ],
-    TestCaseWithoutDecoratorDiagnostic: [
-        "test_logical.py",
-        "test_inplace_assign.py",
-        # TODO: Remove these files from the allow list after it's support PIR mode
-        "test_tensor_hook.py",
-        "test_to_tensor.py",
-        "test_warning.py",
-        "test_typing.py",
-        "test_gradname_parse.py",
-        "test_for_enumerate.py",
-        "test_save_load.py",
-        "test_declarative.py",
-    ],
     TestCaseWithPirApiDecoratorDiagnostic: [],
 }
 
@@ -222,13 +200,6 @@ def visit_ClassDef(self, node: ast.ClassDef):
         self.generic_visit(node)
 
     def check_test_case(self, node: ast.FunctionDef):
-        # Check if the test case has not any decorator
-        if not node.decorator_list:
-            start = Location(node.lineno, node.col_offset)
-            end = Location(node.end_lineno, node.end_col_offset)  # type: ignore
-            self.diagnostics.append(
-                TestCaseWithoutDecoratorDiagnostic(start, end)
-            )
         # Check if the test case use @test_with_pir_api
         for decorator in node.decorator_list:
             decorator_str = ast_to_source_code(decorator).strip()
@@ -347,7 +318,6 @@ def main():
         (
             UseToStaticAsDecoratorDiagnostic,
             TestClassInheritFromTestCaseDiagnostic,
-            TestCaseWithoutDecoratorDiagnostic,
             TestCaseWithPirApiDecoratorDiagnostic,
         ),
     )
 
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import importlib
 import inspect
 import logging
@@ -89,7 +91,7 @@ def lower_case_name(self):
 DEFAULT_TO_STATIC_MODE = (
     ToStaticMode.AST | ToStaticMode.SOT | ToStaticMode.SOT_MGS10
 )
-DEFAULT_IR_MODE = IrMode.LEGACY_IR | IrMode.PT
+DEFAULT_IR_MODE = IrMode.PT | IrMode.PIR
 
 DISABLED_TO_STATIC_TEST_FILES = {
     ToStaticMode.AST: [],
@@ -303,9 +305,9 @@ def __new__(cls, name, bases, attrs):
             for to_static_mode, ir_mode in to_static_with_ir_modes:
                 if (
                     to_static_mode == ToStaticMode.SOT_MGS10
-                    and ir_mode != IrMode.LEGACY_IR
+                    and ir_mode != IrMode.PIR
                 ):
-                    # SOT_MGS10 only test with LEGACY_IR
+                    # SOT_MGS10 only test with PIR
                     continue
                 new_attrs[
                     Dy2StTestMeta.test_case_name(
@@ -347,7 +349,7 @@ def decorator(fn):
     return decorator
 
 
-def disable_test_case(flags):
+def disable_test_case(flags: tuple[ToStaticMode, IrMode]):
     def decorator(fn):
         disabled_test_cases = getattr(fn, "disabled_test_cases", [])
         disabled_test_cases.append(flags)
@@ -389,6 +391,11 @@ def test_legacy_and_pt(fn):
     return fn
 
 
+def test_pt_and_pir(fn):
+    fn = set_ir_mode(IrMode.PT | IrMode.PIR)(fn)
+    return fn
+
+
 def test_legacy_and_pir(fn):
     fn = set_ir_mode(IrMode.LEGACY_IR | IrMode.PIR)(fn)
     return fn
 
@@ -15,16 +15,12 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import (
-    Dy2StTestBase,
-    test_legacy_and_pt_and_pir,
-)
+from dygraph_to_static_utils import Dy2StTestBase
 
 import paddle
 
 
 class TestCpuCuda(Dy2StTestBase):
-    @test_legacy_and_pt_and_pir
     def test_cpu_cuda(self):
         def func(x):
             x = paddle.to_tensor([1, 2, 3, 4])
@@ -40,7 +36,6 @@ def func(x):
 
 
 class TestToTensor(Dy2StTestBase):
-    @test_legacy_and_pt_and_pir
     def test_to_tensor_with_variable_list(self):
         def func(x):
             ones = paddle.to_tensor(1)
@@ -58,7 +53,6 @@ def func(x):
 
 
 class TestToTensor1(Dy2StTestBase):
-    @test_legacy_and_pt_and_pir
     def test_to_tensor_with_variable_list(self):
         def func(x):
             ones = paddle.to_tensor([1])
@@ -76,27 +70,8 @@ def func(x):
             rtol=1e-05,
         )
 
-    @test_legacy_and_pt_and_pir
-    def test_to_tensor_with_variable_list_sot(self):
-        def func(x):
-            ones = paddle.to_tensor([1])
-            twos = paddle.to_tensor([2])
-            """ we ignore the [3] and [4], they will be assign to a variable, and is regard as scalar.
-                TODO: deal with this case after 0-dim tensor is developed.
-            """
-            x = paddle.to_tensor([ones, twos, [3], [4]])
-            return x
-
-        x = paddle.to_tensor([3])
-        np.testing.assert_allclose(
-            paddle.jit.to_static(func)(x),
-            np.array([[1], [2], [3], [4]]),
-            rtol=1e-05,
-        )
-
 
 class TestToTensor2(Dy2StTestBase):
-    @test_legacy_and_pt_and_pir
     def test_to_tensor_with_variable_list(self):
         def func(x):
             x = paddle.to_tensor([[1], [2], [3], [4]])
@@ -109,19 +84,6 @@ def func(x):
             rtol=1e-05,
         )
 
-    @test_legacy_and_pt_and_pir
-    def test_to_tensor_with_variable_list_sot(self):
-        def func(x):
-            x = paddle.to_tensor([[1], [2], [3], [4]])
-            return x
-
-        x = paddle.to_tensor([3])
-        np.testing.assert_allclose(
-            paddle.jit.to_static(func)(x),
-            np.array([[1], [2], [3], [4]]),
-            rtol=1e-05,
-        )
-
 
 if __name__ == '__main__':
     unittest.main()
@@ -18,6 +18,7 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     test_ast_only,
+    test_legacy_and_pt,
     test_legacy_and_pt_and_pir,
 )
 
@@ -77,6 +78,7 @@ def test_dygraph_static_same_result(self):
 
 class TestToStaticInfenrenceTensorRTModel(Dy2StTestBase):
     @test_ast_only
+    @test_legacy_and_pt
     def test_dygraph_static_same_result(self):
         if paddle_infer.get_trt_compile_version()[0] == 0:
             return
@@ -116,6 +118,7 @@ def test_dygraph_static_same_result(self):
 
 class TestToStaticInputListModel(Dy2StTestBase):
     @test_ast_only
+    @test_legacy_and_pt
     def test_dygraph_static_same_result(self):
         hidd = 1024
         batch = 4096
 
@@ -14,7 +14,11 @@
 
 import unittest
 
-from dygraph_to_static_utils import Dy2StTestBase, test_ast_only
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    test_ast_only,
+    test_pt_only,
+)
 
 import paddle
 from paddle.static import InputSpec
@@ -78,6 +82,7 @@ def expected_results(self):
         }
 
     @test_ast_only
+    @test_pt_only
     def test_set_op_attrs(self):
         net = NetWithOpAttr(self.in_num, self.out_num)
         # set attrs
@@ -120,6 +125,7 @@ def check_op_attrs(self, main_program):
                             self.assertEqual(op_val, expect_val)
 
     @test_ast_only
+    @test_pt_only
     def test_set_op_attrs_with_sub_block(self):
         net = NetWithOpAttr(self.in_num, self.out_num)
         # set attrs
@@ -140,6 +146,7 @@ def test_set_op_attrs_with_sub_block(self):
         self.assertEqual(len(net.linear._forward_pre_hooks), 0)
         self.assertEqual(len(net.linear._forward_post_hooks), 0)
 
+    @test_pt_only
     def test_type_error(self):
         net = NetWithOpAttr(self.in_num, self.out_num)
         # attrs should be dict
 
@@ -19,31 +19,12 @@
 import paddle
 from paddle.base import core
 from paddle.jit.dy2static import (
-    partial_program,
     pir_partial_program,
     program_translator,
 )
 
 
 class TestPartiaProgramLayerHook(Dy2StTestBase):
-    # TODO(dev): Remove this after PIR becomes the default.
-    def setUp(self):
-        self._hook = partial_program.PartialProgramLayerHook()
-
-    @test_ast_only
-    def test_before_append_backward(self):
-        self.assertIsNone(self._hook.before_append_backward(None))
-
-    @test_ast_only
-    def test_after_append_backward(self):
-        self.assertIsNone(self._hook.after_append_backward(None, 0))
-
-    @test_ast_only
-    def test_after_infer(self):
-        self.assertIsNone(self._hook.after_infer(None))
-
-
-class TestPirPartiaProgramLayerHook(Dy2StTestBase):
     def setUp(self):
         self._hook = pir_partial_program.PartialProgramLayerHook()
 
@@ -66,45 +47,6 @@ def test_after_infer(self):
 
 
 class TestPrimHook(Dy2StTestBase):
-    def init_test_params(self):
-        core._set_prim_all_enabled(False)
-
-        def f():
-            return paddle.nn.functional.dropout(paddle.rand((1,)))
-
-        concrete_program, partial_program = paddle.jit.to_static(
-            f, full_graph=True
-        ).get_concrete_program()
-        hook = program_translator.PrimHooker(
-            concrete_program.main_program, None
-        )
-        forward = partial_program.forward_program
-        whole = partial_program._train_program
-
-        return hook, forward, whole
-
-    @test_ast_only
-    def test_before_append_backward(self):
-        hook, forward, whole = self.init_test_params()
-        core._set_prim_all_enabled(True)
-        hook.before_append_backward(forward)
-        self.assertNotIn(
-            'dropout', tuple(op.type for op in forward.blocks[0].ops)
-        )
-        core._set_prim_all_enabled(False)
-
-    @test_ast_only
-    def test_after_append_backward(self):
-        hook, forward, whole = self.init_test_params()
-        core._set_prim_all_enabled(True)
-        hook.after_append_backward(whole, 0)
-        self.assertNotIn(
-            'dropout_grad', tuple(op.type for op in whole.blocks[0].ops)
-        )
-        core._set_prim_all_enabled(False)
-
-
-class TestPirPrimHook(Dy2StTestBase):
     def setUp(self):
         core._set_prim_all_enabled(True)
         with paddle.pir_utils.IrGuard():