【PIR API adaptor No.235、236】 Migrate paddle.unbind & unfold into pir

paddle/phi/infermeta/unary.cc

@@ -4809,6 +4809,7 @@ void UnfoldInferMeta(const MetaTensor& x,
   }
   out_dims.push_back(output_col_length);
   out->set_dims(phi::make_ddim(out_dims));
+  out->set_dtype(x.dtype());
 }
 
 void UniformRandomInplaceInferMeta(const MetaTensor& x,

python/paddle/nn/functional/common.py

@@ -109,7 +109,9 @@ def unfold(x, kernel_sizes, strides=1, paddings=0, dilations=1, name=None):
 
     helper = LayerHelper("unfold", **locals())
 
-    check_variable_and_dtype(x, 'x', ['float32', 'float64'], 'unfold')
+    check_variable_and_dtype(
+        x, 'x', ['uint16', 'float16', 'float32', 'float64'], 'unfold'
+    )
 
     assert len(x.shape) == 4, "input should be the format of [N, C, H, W]"
 
@@ -155,7 +157,7 @@ def unfold(x, kernel_sizes, strides=1, paddings=0, dilations=1, name=None):
             "of 2 or 4 integers"
         )
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.unfold(x, kernel_sizes, strides, paddings, dilations)
 
     out = helper.create_variable_for_type_inference(dtype=x.dtype)

python/paddle/tensor/manipulation.py

@@ -2969,7 +2969,7 @@ def unbind(input, axis=0):
             f'The axis must in range({-input.ndim}, {input.ndim}).'
         )
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.unbind(input, axis)
     else:
         if isinstance(axis, np.generic):

test/legacy_test/test_unbind_op.py

@@ -18,29 +18,36 @@
 from op_test import OpTest, convert_float_to_uint16
 
 import paddle
-from paddle import base, tensor
+from paddle import base, static, tensor
 from paddle.base import Program, program_guard
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestUnbind(unittest.TestCase):
+    @test_with_pir_api
     def test_unbind(self):
         paddle.enable_static()
 
-        x_1 = paddle.static.data(shape=[2, 3], dtype='float32', name='x_1')
-        [out_0, out_1] = tensor.unbind(input=x_1, axis=0)
-        input_1 = np.random.random([2, 3]).astype("float32")
-        axis = paddle.static.data(shape=[], dtype='int32', name='axis')
-        exe = base.Executor(place=base.CPUPlace())
-
-        [res_1, res_2] = exe.run(
-            base.default_main_program(),
-            feed={"x_1": input_1, "axis": 0},
-            fetch_list=[out_0, out_1],
-        )
-
-        np.testing.assert_array_equal(res_1, input_1[0, 0:100])
-        np.testing.assert_array_equal(res_2, input_1[1, 0:100])
-
+        main_program = static.Program()
+        startup_program = static.Program()
+        with static.program_guard(
+            main_program=main_program, startup_program=startup_program
+        ):
+            x_1 = paddle.static.data(shape=[2, 3], dtype='float32', name='x_1')
+            [out_0, out_1] = tensor.unbind(input=x_1, axis=0)
+            input_1 = np.random.random([2, 3]).astype("float32")
+            axis = paddle.static.data(shape=[], dtype='int32', name='axis')
+            exe = base.Executor(place=base.CPUPlace())
+
+            [res_1, res_2] = exe.run(
+                feed={"x_1": input_1, "axis": 0},
+                fetch_list=[out_0, out_1],
+            )
+
+            np.testing.assert_array_equal(res_1, input_1[0, 0:100])
+            np.testing.assert_array_equal(res_2, input_1[1, 0:100])
+
+    @test_with_pir_api
     def test_unbind_static_fp16_gpu(self):
         if paddle.base.core.is_compiled_with_cuda():
             place = paddle.CUDAPlace(0)
@@ -81,6 +88,7 @@ def test_unbind_dygraph(self):
 
 
 class TestLayersUnbind(unittest.TestCase):
+    @test_with_pir_api
     def test_layers_unbind(self):
         paddle.enable_static()
 
@@ -91,7 +99,6 @@ def test_layers_unbind(self):
         exe = base.Executor(place=base.CPUPlace())
 
         [res_1, res_2] = exe.run(
-            base.default_main_program(),
             feed={"x_1": input_1, "axis": 0},
             fetch_list=[out_0, out_1],
         )
@@ -137,10 +144,10 @@ def _set_op_type(self):
         self.op_type = "unbind"
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], ['out0', 'out1', 'out2'])
+        self.check_grad(['X'], ['out0', 'out1', 'out2'], check_pir=True)
 
 
 class TestUnbindOp1(TestUnbindOp):
@@ -149,7 +156,7 @@ def initParameters(self):
         self.num = 2
 
     def test_check_grad(self):
-        self.check_grad(['X'], ['out0', 'out1'])
+        self.check_grad(['X'], ['out0', 'out1'], check_pir=True)
 
     def outReshape(self):
         self.out[0] = self.out[0].reshape((3, 2))
@@ -162,7 +169,7 @@ def initParameters(self):
         self.num = 2
 
     def test_check_grad(self):
-        self.check_grad(['X'], ['out0', 'out1'])
+        self.check_grad(['X'], ['out0', 'out1'], check_pir=True)
 
     def outReshape(self):
         self.out[0] = self.out[0].reshape((3, 2))
@@ -178,7 +185,7 @@ def setAxis(self):
         self.attrs = {'axis': -1}
 
     def test_check_grad(self):
-        self.check_grad(['X'], ['out0', 'out1'])
+        self.check_grad(['X'], ['out0', 'out1'], check_pir=True)
 
     def outReshape(self):
         self.out[0] = self.out[0].reshape((3, 2))
@@ -194,7 +201,7 @@ def setAxis(self):
         self.attrs = {'axis': -2}
 
     def test_check_grad(self):
-        self.check_grad(['X'], ['out0', 'out1'])
+        self.check_grad(['X'], ['out0', 'out1'], check_pir=True)
 
     def outReshape(self):
         self.out[0] = self.out[0].reshape((3, 2))
@@ -228,7 +235,7 @@ def get_dtype(self):
         return np.float16
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
 
 class TestUnbindBF16Op(OpTest):
@@ -264,13 +271,14 @@ def _set_op_type(self):
         self.op_type = "unbind"
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         pass
 
 
 class TestUnbindAxisError(unittest.TestCase):
+    @test_with_pir_api
     def test_errors(self):
         with program_guard(Program(), Program()):
             x = paddle.static.data(shape=[2, 3], dtype='float32', name='x')

test/legacy_test/test_unfold_op.py

@@ -139,7 +139,7 @@ def init_dtype(self):
         self.dtype = np.float64
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
         self.check_grad(['X'], 'Y')
@@ -199,10 +199,10 @@ def setUp(self):
         self.place = core.CUDAPlace(0)
 
     def test_check_output(self):
-        self.check_output_with_place(self.place)
+        self.check_output_with_place(self.place, check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad_with_place(self.place, ['X'], 'Y')
+        self.check_grad_with_place(self.place, ['X'], 'Y', check_pir=True)
 
 
 class TestUnfoldAPI(TestUnfoldOp):