diff --git a/python/paddle/fluid/tests/unittests/test_diag.py b/python/paddle/fluid/tests/unittests/test_diag.py index b270d47316f40..0e645e5fb7d65 100644 --- a/python/paddle/fluid/tests/unittests/test_diag.py +++ b/python/paddle/fluid/tests/unittests/test_diag.py @@ -18,6 +18,7 @@ from op_test import OpTest import paddle +import paddle.fluid.core as core from paddle.fluid import Program, program_guard @@ -42,14 +43,26 @@ def init_config(self): self.case = np.array([3], dtype='int32') +class TestDiagOpFp16(unittest.TestCase): + def test_fp16(self): + x_np = np.array([3], dtype='float16') + with paddle.static.program_guard(paddle.static.Program()): + x = paddle.static.data(shape=[1, 0], name='x', dtype='float16') + out = paddle.diag(x) + if core.is_compiled_with_cuda(): + place = paddle.CUDAPlace(0) + exe = paddle.static.Executor(place) + exe.run(paddle.static.default_startup_program()) + out = exe.run(feed={'x': x_np}, fetch_list=[out]) + + class TestDiagError(unittest.TestCase): def test_errors(self): paddle.enable_static() with program_guard(Program(), Program()): def test_diag_type(): - x = [1, 2, 3] - output = paddle.diag(x=x) + return paddle.diag(x=[1, 2, 3]) self.assertRaises(TypeError, test_diag_type) diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py index 83379ff69c8a4..67e63d5c29434 100644 --- a/python/paddle/tensor/creation.py +++ b/python/paddle/tensor/creation.py @@ -1627,7 +1627,7 @@ def diag(x, offset=0, padding_value=0, name=None): If ``offset`` < 0, it is subdiagonal. Args: - x (Tensor): The input tensor. Its shape is either 1-D or 2-D. Its data type should be float32, float64, int32, int64. + x (Tensor): The input tensor. Its shape is either 1-D or 2-D. 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. padding_value (int|float, optional): Use this value to fill the area outside the specified diagonal band. Only takes effect when the input is a 1-D Tensor. The default value is 0. name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None. @@ -1694,7 +1694,7 @@ def diag(x, offset=0, padding_value=0, name=None): check_dtype( x.dtype, 'x', - ['float32', 'float64', 'int32', 'int64'], + ['float16', 'float32', 'float64', 'int32', 'int64'], 'diag_v2', ) check_type(offset, 'offset', (int), 'diag_v2')