[PIR] A-15 Adapt split test_errors (#61039)

enkilee · web-flow · commit 8651551d9069 · 2024-02-22T10:05:59.000+08:00
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -2512,6 +2512,12 @@ def split(x, num_or_sections, axis=0, name=None):
             dim = (len(input.shape) + dim) if dim < 0 else dim
 
         input_shape = input.shape
+
+        if not isinstance(num_or_sections, (int, list, tuple)):
+            raise TypeError(
+                "The type of 'num_or_sections' in split must be int, list or tuple in imperative mode, but "
+                "received %s." % (type(num_or_sections))
+            )
         if isinstance(num_or_sections, int):
             assert num_or_sections > 0, 'num_or_sections must be than 0.'
             if isinstance(dim, int) and input_shape[dim] > 0:
diff --git a/test/legacy_test/test_split_op.py b/test/legacy_test/test_split_op.py
@@ -19,7 +19,8 @@
 
 import paddle
 from paddle import base
-from paddle.base import Program, core, program_guard
+from paddle.base import core
+from paddle.framework import in_pir_mode
 from paddle.pir_utils import test_with_pir_api
 
 
@@ -361,57 +362,64 @@ def test_api(self):
             np.testing.assert_array_equal(res_5, out[2])
 
 
-class TestSplitOpError(unittest.TestCase):
-    def test_errors(self):
+class TestSplitOpErrorStatic(unittest.TestCase):
+    @test_with_pir_api
+    def test_errors_with_static(self):
         paddle.enable_static()
-        with program_guard(Program(), Program()):
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
             # The type of axis in split_op should be int or Variable.
             def test_axis_type():
-                x6 = paddle.static.data(
-                    shape=[-1, 4], dtype='float16', name='x3'
+                x5 = paddle.static.data(
+                    shape=[-1, 4], dtype='float16', name='x5'
                 )
-                paddle.split(x=x6, num_or_sections=2, axis=3.2)
+                paddle.split(x=x5, num_or_sections=2, axis=3.2)
 
             self.assertRaises(TypeError, test_axis_type)
 
-            # The type of axis in split_op should be int or Variable.
-            def test_axis_variable_type():
-                x9 = paddle.static.data(
-                    shape=[-1, 4], dtype='float16', name='x9'
-                )
-                x10 = paddle.static.data(
-                    shape=[-1, 1], dtype='float16', name='x10'
-                )
-                paddle.split(x=x9, num_or_sections=2, axis=x10)
+            if not in_pir_mode():
+                # The type of axis in split_op should be int or Variable.
+                def test_axis_variable_type():
+                    x9 = paddle.static.data(
+                        shape=[-1, 4], dtype='float16', name='x9'
+                    )
+                    x10 = paddle.static.data(
+                        shape=[-1, 1], dtype='float16', name='x10'
+                    )
+                    paddle.split(x=x9, num_or_sections=2, axis=x10)
 
-            self.assertRaises(TypeError, test_axis_variable_type)
+                self.assertRaises(TypeError, test_axis_variable_type)
 
             # The type of num_or_sections in split_op should be int, tuple or list.
             def test_num_or_sections_type():
                 x6 = paddle.static.data(
-                    shape=[-1, 4], dtype='float16', name='x4'
+                    shape=[-1, 4], dtype='float16', name='x6'
                 )
                 paddle.split(x=x6, num_or_sections=2.1, axis=3)
 
             self.assertRaises(TypeError, test_num_or_sections_type)
 
             def test_num_or_sections_type_tensor():
                 x7 = paddle.static.data(
-                    shape=[-1, 4], dtype='float16', name='x5'
+                    shape=[-1, 4], dtype='float16', name='x7'
                 )
                 paddle.split(input=x7, num_or_sections=2.1, dim=3)
 
             self.assertRaises(TypeError, test_num_or_sections_type_tensor)
 
             def test_axis_type_tensor():
                 x8 = paddle.static.data(
-                    shape=[-1, 4], dtype='float16', name='x6'
+                    shape=[-1, 4], dtype='float16', name='x8'
                 )
                 paddle.split(input=x8, num_or_sections=2, dim=3.2)
 
             self.assertRaises(TypeError, test_axis_type_tensor)
         paddle.disable_static()
 
+
+class TestSplitOpErrorDynamic(unittest.TestCase):
+    def test_errors_with_dynamic(self):
         with paddle.base.dygraph.guard():
 
             def test_0_num_tensor():
