support lazy_init in pir, but not add test case (#67379)

wanghuancoder · web-flow · commit 97788c9c9762 · 2024-08-14T14:05:02.000+08:00
diff --git a/python/paddle/nn/initializer/initializer.py b/python/paddle/nn/initializer/initializer.py
@@ -25,6 +25,7 @@
     EagerParamBase,
     default_main_program,
     in_dygraph_mode,
+    use_pir_api,
 )
 from .lazy_init import lazy_init_helper
 
@@ -86,7 +87,10 @@ def _lazy_init(
         def init_op_creator(
             forward, param: paddle.Tensor, block: paddle.pir.Block | None
         ):
-            new_var = param._to_static_var(True, block=block)
+            if use_pir_api():
+                new_var = param
+            else:
+                new_var = param._to_static_var(True, block=block)
             # Record initializer operator
             with lazy_init_helper():
                 forward(new_var, block)
diff --git a/python/paddle/nn/layer/layers.py b/python/paddle/nn/layer/layers.py
@@ -2484,11 +2484,15 @@ def _startup_program(self) -> Program:
 
         NOTE(dev): This is a very low level API and only for inner developer.
         """
-        startup_program = Program()
-        for param in self.parameters():
-            param._create_init_op(startup_program.global_block())
-
-        return startup_program
+        startup_program = paddle.base.Program()
+        main_program = paddle.base.Program()
+        with paddle.base.program_guard(main_program, startup_program):
+            for param in self.parameters():
+                param._create_init_op(startup_program.global_block())
+        if paddle.framework.use_pir_api():
+            return main_program
+        else:
+            return startup_program
 
     # [aliases] Compatible with old method names
     set_dict = set_state_dict
diff --git a/test/legacy_test/test_lazy_init.py b/test/legacy_test/test_lazy_init.py
@@ -57,13 +57,17 @@ def clear_nameset(self):
         unique_name.dygraph_parameter_name_checker._name_set = set()
 
     def test_wrapper(self):
-        with LazyGuard():
-            fc = Linear(
-                10, 10, weight_attr=self.weight_attr, bias_attr=self.bias_attr
-            )
-        program = fc._startup_program()
-        print(program)
-        self.check_program(program)
+        with paddle.pir_utils.OldIrGuard():
+            paddle.disable_static()
+            with LazyGuard():
+                fc = Linear(
+                    10,
+                    10,
+                    weight_attr=self.weight_attr,
+                    bias_attr=self.bias_attr,
+                )
+            program = fc._startup_program()
+            self.check_program(program)
 
     def check_program(self, program):
         self.assertEqual(program.block(0).var("weight").shape, (10, 10))
@@ -115,14 +119,19 @@ def forward(self, x):
 
 class TestNestModelLazy(TestInitializerBase):
     def test_wrapper(self):
-        with LazyGuard():
-            base_model = Linear(
-                10, 10, weight_attr=self.weight_attr, bias_attr=self.bias_attr
-            )
-            nest_model = NestModel(base_model)
-
-        self.check_data(nest_model)
-        self.check_program(nest_model)
+        with paddle.pir_utils.OldIrGuard():
+            paddle.disable_static()
+            with LazyGuard():
+                base_model = Linear(
+                    10,
+                    10,
+                    weight_attr=self.weight_attr,
+                    bias_attr=self.bias_attr,
+                )
+                nest_model = NestModel(base_model)
+
+            self.check_data(nest_model)
+            self.check_program(nest_model)
 
     def check_data(self, model):
         x = paddle.randn([2, 10])
