[fp16] suppot fp16 in diagflat (#50945)

PaddlePaddle · Mar 1, 2023 · af149c0 · af149c0
1 parent 48060b2
commit af149c0
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/python/paddle/fluid/tests/unittests/test_diagflat.py b/python/paddle/fluid/tests/unittests/test_diagflat.py
@@ -104,6 +104,29 @@ def test_gpu(self):
         with paddle.static.program_guard(Program()):
             self.run_static(use_gpu=True)
 
+    def test_fp16_with_gpu(self, use_gpu=False):
+        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([10, 10]).astype("float16")
+                x = paddle.static.data(
+                    name="x", shape=[10, 10], dtype="float16"
+                )
+
+                y = paddle.diagflat(x)
+                expected = np.diagflat(input)
+
+                exe = paddle.static.Executor(place)
+                res = exe.run(
+                    paddle.static.default_main_program(),
+                    feed={
+                        "x": input,
+                    },
+                    fetch_list=[y],
+                )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -1498,7 +1498,7 @@ def diagflat(x, offset=0, name=None):
     If ``offset`` < 0, it is subdiagonal.
 
     Args:
-        x (Tensor): The input tensor. It can be any shape. Its data type should be float32, float64, int32, int64.
+        x (Tensor): The input tensor. It can be any shape. Its data type should be float16, float32, float64, int32, int64.
         offset (int, optional): The diagonal offset. A positive value represents superdiagonal, 0 represents the main diagonal, and a negative value represents subdiagonal. Default: 0 (main diagonal).
         name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None.
 
@@ -1577,7 +1577,10 @@ def diagflat(x, offset=0, name=None):
         padding_value = 0
         check_type(x, 'x', (Variable), 'diagflat')
         check_dtype(
-            x.dtype, 'x', ['float32', 'float64', 'int32', 'int64'], 'diagflat'
+            x.dtype,
+            'x',
+            ['float16', 'float32', 'float64', 'int32', 'int64'],
+            'diagflat',
         )
         check_type(offset, 'offset', (int), 'diagflat')