test_rot90_passed (#60193)

changeyoung98 · web-flow · commit 95d7605721d4 · 2023-12-21T11:16:01.000+08:00
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -1580,7 +1580,7 @@ def rot90(x, k=1, axes=[0, 1], name=None):
     """
 
     helper = LayerHelper("rot90", **locals())
-    check_type(x, 'X', (Variable), 'rot90')
+    check_type(x, 'X', (Variable, paddle.pir.OpResult), 'rot90')
     dtype = helper.input_dtype('x')
     check_dtype(
         dtype,
diff --git a/test/legacy_test/test_rot90_op.py b/test/legacy_test/test_rot90_op.py
@@ -18,11 +18,13 @@
 
 import paddle
 from paddle import base
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestRot90_API(unittest.TestCase):
     """Test rot90 api."""
 
+    @test_with_pir_api
     def test_static_graph(self):
         paddle.enable_static()
         startup_program = base.Program()
@@ -53,6 +55,7 @@ def test_static_graph(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_k_0(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -82,6 +85,7 @@ def test_static_k_0(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_k_2(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -111,6 +115,7 @@ def test_static_k_2(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_k_3(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -140,6 +145,7 @@ def test_static_k_3(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_neg_k_1(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -169,6 +175,7 @@ def test_static_neg_k_1(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_neg_k_2(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -198,6 +205,7 @@ def test_static_neg_k_2(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_neg_k_3(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -227,6 +235,7 @@ def test_static_neg_k_3(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_static_neg_k_4(self):
         paddle.enable_static()
         input = paddle.static.data(name='input', dtype='float32', shape=[2, 3])
@@ -256,6 +265,7 @@ def test_static_neg_k_4(self):
                 msg='rot90 output is wrong, out =' + str(out_np),
             )
 
+    @test_with_pir_api
     def test_error_api(self):
         paddle.enable_static()
 
