[XPU] fix unbind op ut (#68010)

houj04 · web-flow · commit c4e2eec9655a · 2024-09-06T10:02:41.000+08:00
* [XPU] revert unbind op ut

* fix ut.
diff --git a/test/xpu/test_unbind_op_xpu.py b/test/xpu/test_unbind_op_xpu.py
@@ -20,37 +20,50 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_uint16_to_float
 from op_test_xpu import XPUOpTest
 
 import paddle
+from paddle import base, tensor
+from paddle.base import Program, program_guard
 
 paddle.enable_static()
 
-"""
-class TestUnbind(unittest.TestCase):
-        def test_unbind(self):
-            self.dtype = self.in_type
-            self.place = paddle.XPUPlace(0)
-            x_1 = paddle.static.data(shape=[2, 3], dtype=self.dtype, name='x_1')
-            [out_0, out_1] = tensor.unbind(input=x_1, axis=0)
-            input_1 = np.random.random([2, 3]).astype(self.dtype)
-            axis = paddle.static.data(shape=[], dtype='int32', name='axis')
-            exe = base.Executor(place=self.place)
 
-            [res_1, res_2] = exe.run(
-                base.default_main_program(),
-                feed={"x_1": input_1, "axis": 0},
-                fetch_list=[out_0, out_1],
-            )
+class XPUTestUnbindOP(XPUOpTestWrapper):
+    def __init__(self):
+        self.op_name = 'unbind'
+        self.use_dynamic_create_class = False
 
-            np.testing.assert_array_equal(res_1, input_1[0, 0:100])
-            np.testing.assert_array_equal(res_2, input_1[1, 0:100])
+    class TestUnbind(unittest.TestCase):
+        def test_unbind(self):
+            with program_guard(Program(), Program()):
+                self.dtype = self.in_type
+                self.place = paddle.XPUPlace(0)
+                x_1 = paddle.static.data(
+                    shape=[2, 3], dtype=self.dtype, name='x_1'
+                )
+                [out_0, out_1] = tensor.unbind(input=x_1, axis=0)
+                input_1 = np.random.random([2, 3]).astype(self.dtype)
+                axis = paddle.static.data(shape=[], dtype='int32', name='axis')
+                exe = base.Executor(place=self.place)
+
+                [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])
 
         def test_unbind_dygraph(self):
             with base.dygraph.guard():
                 self.dtype = self.in_type
                 self.place = paddle.XPUPlace(0)
                 np_x = np.random.random([2, 3]).astype(self.dtype)
+                if self.dtype == np.uint16:
+                    np_x = convert_uint16_to_float(np_x)
                 x = paddle.to_tensor(np_x)
                 x.stop_gradient = False
                 [res_1, res_2] = paddle.unbind(x, 0)
@@ -68,29 +81,25 @@ def test_unbind_dygraph_final_state(self):
 
     class TestLayersUnbind(unittest.TestCase):
         def test_layers_unbind(self):
-            self.dtype = self.in_type
-            self.place = paddle.XPUPlace(0)
-            x_1 = paddle.static.data(shape=[2, 3], dtype=self.dtype, name='x_1')
-            [out_0, out_1] = paddle.unbind(input=x_1, axis=0)
-            input_1 = np.random.random([2, 3]).astype(self.dtype)
-            axis = paddle.static.data(shape=[], dtype='int32', name='axis')
-            exe = base.Executor(place=self.place)
-
-            [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])
-"""
-
-
-class XPUTestUnbindOP(XPUOpTestWrapper):
-    def __init__(self):
-        self.op_name = 'unbind'
-        self.use_dynamic_create_class = False
+            with program_guard(Program(), Program()):
+                self.dtype = self.in_type
+                self.place = paddle.XPUPlace(0)
+                x_1 = paddle.static.data(
+                    shape=[2, 3], dtype=self.dtype, name='x_1'
+                )
+                [out_0, out_1] = paddle.unbind(input=x_1, axis=0)
+                input_1 = np.random.random([2, 3]).astype(self.dtype)
+                axis = paddle.static.data(shape=[], dtype='int32', name='axis')
+                exe = base.Executor(place=self.place)
+
+                [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])
 
     class TestUnbindOp(XPUOpTest):
         def initParameters(self):
@@ -186,8 +195,6 @@ def outReshape(self):
             self.out[0] = self.out[0].reshape((3, 2))
             self.out[1] = self.out[1].reshape((3, 2))
 
-
-"""
     class TestUnbindAxisError(unittest.TestCase):
         def test_errors(self):
             with program_guard(Program(), Program()):
@@ -204,7 +211,7 @@ def test_invalid_axis():
                     tensor.unbind(input=x, axis=2)
 
                 self.assertRaises(ValueError, test_invalid_axis)
-"""
+
 
 support_types = get_xpu_op_support_types('unbind')
 for stype in support_types:
