【PIR API adaptor No.14-15】Assign and Bilinear

python/paddle/nn/initializer/Bilinear.py

@@ -17,9 +17,10 @@
 from paddle import _C_ops
 
 from ...base import core, framework, unique_name
-from ...base.framework import _current_expected_place, in_dygraph_mode
+from ...base.framework import _current_expected_place, in_dynamic_or_pir_mode
 from .initializer import Initializer
 
+
 __all__ = []
 
 
@@ -140,7 +141,7 @@ def forward(self, var, block=None):
         if np.prod(shape) > 1024 * 1024:
             raise ValueError("The size of input is too big. ")
 
-        if in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             _C_ops.assign_value_(
                 out_var,
                 list(shape),

python/paddle/nn/initializer/assign.py

@@ -16,7 +16,7 @@
 
 from ...base import core, framework, unique_name
 from ...base.data_feeder import check_type
-from ...base.framework import _current_expected_place, in_dygraph_mode
+from ...base.framework import _current_expected_place, in_dynamic_or_pir_mode
 from .initializer import Initializer
 
 __all__ = []
@@ -98,7 +98,7 @@ def forward(self, var, block=None):
                 "saving it to file and 'load_op' to load it"
             )
 
-        if in_dygraph_mode():
+        if in_dynamic_or_pir_mode():
             _C_ops.assign_value_(
                 out_var,
                 list(self._value.shape),

test/legacy_test/test_initializer.py

@@ -23,6 +23,7 @@
 from paddle import base
 from paddle.base import framework
 from paddle.base.core import VarDesc
+from paddle.pir_utils import test_with_pir_api
 from paddle.regularizer import L2Decay
 
 DELTA = 0.00001
@@ -541,6 +542,7 @@ def test_msra_initializer_bf16(self):
 
 
 class TestBilinearInitializer(unittest.TestCase):
+    @test_with_pir_api
     def test_bilinear_initializer(self, dtype="float32"):
         """Test the bilinear initializer with supplied arguments"""
         program = framework.Program()
@@ -628,6 +630,7 @@ def test_bilinear_initializer_fp16(self):
 
 
 class TestNumpyArrayInitializer(unittest.TestCase):
+    @test_with_pir_api
     def test_numpy_array_initializer(self, dtype="float32"):
         """Test the numpy array initializer with supplied arguments"""
         import numpy
@@ -876,6 +879,7 @@ def run_dynamic_graph(dtype):
                 )
             return w
 
+        @test_with_pir_api
         def run_static_graph(dtype):
             with static_guard():
                 exe = paddle.static.Executor(paddle.CPUPlace())
@@ -914,6 +918,7 @@ def run_dynamic_graph(dtype):
                 )
             return w
 
+        @test_with_pir_api
         def run_static_graph(dtype):
             with static_guard():
                 exe = paddle.static.Executor(paddle.CPUPlace())