From 5f60b5974426cd4b59e815defb30760ea446891a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E6=98=A5=E4=B9=94?=
 <83450930+Liyulingyue@users.noreply.github.com>
Date: Mon, 27 Feb 2023 15:45:32 +0800
Subject: [PATCH] support fp16 on unbind (#50916)

---
 .../fluid/tests/unittests/test_unbind_op.py   | 23 +++++++++++++++++++
 python/paddle/tensor/manipulation.py          |  7 ++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/python/paddle/fluid/tests/unittests/test_unbind_op.py b/python/paddle/fluid/tests/unittests/test_unbind_op.py
index 431f807efd2..cf1beb5bc87 100644
--- a/python/paddle/fluid/tests/unittests/test_unbind_op.py
+++ b/python/paddle/fluid/tests/unittests/test_unbind_op.py
@@ -42,6 +42,29 @@ class TestUnbind(unittest.TestCase):
         assert np.array_equal(res_1, input_1[0, 0:100])
         assert np.array_equal(res_2, input_1[1, 0:100])
 
+    def test_unbind_static_fp16_gpu(self):
+        if paddle.fluid.core.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                input = np.random.random([2, 3]).astype("float16")
+
+                x = paddle.static.data(name="x", shape=[2, 3], dtype="float16")
+                y = paddle.unbind(x)
+
+                exe = paddle.static.Executor(place)
+                res = exe.run(
+                    paddle.static.default_main_program(),
+                    feed={
+                        "x": input,
+                    },
+                    fetch_list=[y],
+                )
+
+                assert np.array_equal(res[0], input[0, :])
+                assert np.array_equal(res[1], input[1, :])
+
     def test_unbind_dygraph(self):
         with fluid.dygraph.guard():
             np_x = np.random.random([2, 3]).astype("float32")
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
index c8147b1cbe2..2b4caeff7e5 100644
--- a/python/paddle/tensor/manipulation.py
+++ b/python/paddle/tensor/manipulation.py
@@ -2738,7 +2738,7 @@ def unbind(input, axis=0):
     Removes a tensor dimension, then split the input tensor into multiple sub-Tensors.
 
     Args:
-        input (Tensor): The input variable which is an N-D Tensor, data type being float32, float64, int32 or int64.
+        input (Tensor): The input variable which is an N-D Tensor, data type being float16, float32, float64, int32 or int64.
         axis (int32|int64, optional): A scalar with type ``int32|int64`` shape [1]. The dimension along which to unbind.
             If :math:`axis < 0`, the dimension to unbind along is :math:`rank(input) + axis`. Default is 0.
     Returns:
@@ -2785,7 +2785,10 @@ def unbind(input, axis=0):
         check_type(input, 'input', (Variable), 'unbind')
         dtype = helper.input_dtype()
         check_dtype(
-            dtype, 'unbind', ['float32', 'float64', 'int32', 'int64'], 'unbind'
+            dtype,
+            'unbind',
+            ['float16', 'float32', 'float64', 'int32', 'int64'],
+            'unbind',
         )
         outs = [
             helper.create_variable_for_type_inference(
-- 
GitLab