【PIR API adaptor No.314】 Migrate vander into pir (#59573)

DrRyanHuang · web-flow · commit afd668522272 · 2023-12-04T10:52:24.000+08:00
diff --git a/python/paddle/base/variable_index.py b/python/paddle/base/variable_index.py
@@ -415,6 +415,7 @@ def _setitem_static(x, indices, values):
         indices(int|slice|None|Tensor|List|Tuple...): Indices, used to indicate the position of the element to be fetched.
         values(Tensor|Number|Ndarray): values to be assigned to the x.
     """
+    from . import in_dynamic_or_pir_mode
     from .framework import Variable, default_main_program
 
     if x.type == paddle.base.core.VarDesc.VarType.LOD_TENSOR_ARRAY:
@@ -479,7 +480,7 @@ def _setitem_static(x, indices, values):
             attrs["values"] = values
             attrs["shape"] = shape
 
-        elif isinstance(values, Variable):
+        elif isinstance(values, (Variable, paddle.pir.Value)):
             values = values.astype(dtype)
             inputs["ValueTensor"] = values
             value_tensor = values
@@ -491,7 +492,7 @@ def _setitem_static(x, indices, values):
             )
 
         # step3.1: Only basic indexing, use OP set_value to set value.
-        if paddle.in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             if value_tensor is None:
                 return paddle._C_ops.set_value_(
                     x,
@@ -566,13 +567,13 @@ def _setitem_static(x, indices, values):
             _,
             _,
         ) = deal_advanced_index(sub_tensor, advanced_index, True)
-        if not isinstance(values, Variable):
+        if not isinstance(values, (Variable, paddle.pir.Value)):
             values = paddle.assign(values).astype(transed_sub_tensor.dtype)
 
         if values.dtype != transed_sub_tensor.dtype:
             values = values.astype(transed_sub_tensor.dtype)
 
-        if paddle.in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             # NOTE(zoooo0820): directly return result instead of another set_value, after backward bug fixed.
             transed_sub_tensor = transed_sub_tensor.index_put_(
                 adjusted_advanced_index, values
@@ -585,7 +586,7 @@ def _setitem_static(x, indices, values):
         transback_sub_tensor = transed_sub_tensor.transpose(transback_dim)
         inputs["ValueTensor"] = transback_sub_tensor
 
-        if paddle.in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             x._bump_inplace_version()
             output = x
         else:
@@ -610,7 +611,7 @@ def _setitem_static(x, indices, values):
             inplace_map={"Input": "Out"},
         )
 
-        if not paddle.in_dynamic_mode():
+        if not in_dynamic_or_pir_mode():
             # map var to the new output
             paddle.jit.api.ProgramTranslator.get_instance()._inplace_map.add(
                 cur_block.program, x.desc.id(), output
diff --git a/test/legacy_test/test_vander.py b/test/legacy_test/test_vander.py
@@ -18,6 +18,7 @@
 
 import paddle
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 np.random.seed(10)
 
@@ -41,15 +42,22 @@ def setUp(self):
     def api_case(self, N=None, increasing=False):
         paddle.enable_static()
         out_ref = ref_vander(self.x, N, increasing)
-        with paddle.static.program_guard(paddle.static.Program()):
-            x = paddle.static.data('X', self.shape)
-            out = paddle.vander(x, N, increasing)
-            exe = paddle.static.Executor(self.place)
-            res = exe.run(feed={'X': self.x}, fetch_list=[out])
-        if N != 0:
-            np.testing.assert_allclose(res[0], out_ref, rtol=1e-05)
-        else:
-            np.testing.assert_allclose(res[0].size, out_ref.size, rtol=1e-05)
+
+        @test_with_pir_api
+        def test_static_or_pir_mode():
+            with paddle.static.program_guard(paddle.static.Program()):
+                x = paddle.static.data('X', self.shape)
+                out = paddle.vander(x, N, increasing)
+                exe = paddle.static.Executor(self.place)
+                res = exe.run(feed={'X': self.x}, fetch_list=[out])
+            if N != 0:
+                np.testing.assert_allclose(res[0], out_ref, rtol=1e-05)
+            else:
+                np.testing.assert_allclose(
+                    res[0].size, out_ref.size, rtol=1e-05
+                )
+
+        test_static_or_pir_mode()
 
         paddle.disable_static(self.place)
         x = paddle.to_tensor(self.x)
@@ -86,6 +94,7 @@ def test_api_case(N, increasing=False):
         test_api_case(N, increasing=True)
         paddle.enable_static()
 
+    @test_with_pir_api
     def test_errors(self):
         paddle.enable_static()
         with paddle.static.program_guard(paddle.static.Program()):
