diff --git a/python/paddle/fluid/tests/unittests/test_nansum_api.py b/python/paddle/fluid/tests/unittests/test_nansum_api.py index 5c56eb21947ab7ea9856dde66bbb8e7ceab0fdc9..87b05b4245d22863a07a7e55dab321be2ebcf844 100644 --- a/python/paddle/fluid/tests/unittests/test_nansum_api.py +++ b/python/paddle/fluid/tests/unittests/test_nansum_api.py @@ -72,22 +72,58 @@ class API_Test_Nansum(unittest.TestCase): msg='nansum output is wrong, out =' + str(out4_np), ) - def test_error_api(self): + # test nansum api with float16 + def test_static_graph_fp16(self): paddle.enable_static() - - # input dtype error - def run1(): - input = fluid.data(name='input', dtype='float16', shape=[2, 3]) - output = paddle.nansum(input) - - self.assertRaises(TypeError, run1) - - # axis type error - def run2(): - input = fluid.data(name='input', dtype='float16', shape=[2, 3]) - output = paddle.nansum(input, axis=1.2) - - self.assertRaises(TypeError, run2) + startup_program = paddle.static.Program() + train_program = paddle.static.Program() + with paddle.static.program_guard(train_program, startup_program): + input = paddle.static.data( + name='input', dtype='float16', shape=[2, 4] + ) + out1 = paddle.nansum(input) + out2 = paddle.nansum(input, axis=0) + out3 = paddle.nansum(input, axis=-1) + out4 = paddle.nansum(input, axis=1, keepdim=True) + if fluid.core.is_compiled_with_cuda(): + place = paddle.CUDAPlace(0) + exe = paddle.static.Executor(place) + exe.run(startup_program) + + x = np.array( + [[float('nan'), 3, 5, 9], [1, 2, float('-nan'), 7]] + ).astype(np.float16) + res = exe.run( + train_program, + feed={'input': x}, + fetch_list=[out1, out2, out3, out4], + ) + + out1_np = np.array(res[0]) + out2_np = np.array(res[1]) + out3_np = np.array(res[2]) + out4_np = np.array(res[3]) + out1_ref = np.array([27]).astype(np.float16) + out2_ref = np.array([1, 5, 5, 16]).astype(np.float16) + out3_ref = np.array([17, 10]).astype(np.float16) + out4_ref = np.array([[17], [10]]).astype(np.float16) + + self.assertTrue( + (out1_np == out1_ref).all(), + msg='nansum output is wrong, out =' + str(out1_np), + ) + self.assertTrue( + (out2_np == out2_ref).all(), + msg='nansum output is wrong, out =' + str(out2_np), + ) + self.assertTrue( + (out3_np == out3_ref).all(), + msg='nansum output is wrong, out =' + str(out3_np), + ) + self.assertTrue( + (out4_np == out4_ref).all(), + msg='nansum output is wrong, out =' + str(out4_np), + ) def test_dygraph(self): x = np.array( diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py index 8e7828406e2e54f7e93756519a23e1ce13e173ac..591b885439517ed2eaf9edc2b6c82710dcca2a9a 100644 --- a/python/paddle/tensor/math.py +++ b/python/paddle/tensor/math.py @@ -1349,7 +1349,7 @@ def nansum(x, axis=None, dtype=None, keepdim=False, name=None): Computes the sum of tensor elements over the given axis, treating Not a Numbers (NaNs) as zero. Args: - x (Tensor): An N-D Tensor, the data type is float32, float64, int32 or int64. + x (Tensor): An N-D Tensor, the data type is float16, float32, float64, int32 or int64. axis (int|list|tuple, optional): The dimensions along which the nansum is performed. If :attr:`None`, nansum all elements of :attr:`x` and return a Tensor with a single element, otherwise must be in the @@ -1392,7 +1392,7 @@ def nansum(x, axis=None, dtype=None, keepdim=False, name=None): out6 = paddle.nansum(y, axis=[0, 1]) # [9, 18] """ check_variable_and_dtype( - x, 'x', ['float32', 'float64', 'int32', 'int64'], 'nansum' + x, 'x', ['float16', 'float32', 'float64', 'int32', 'int64'], 'nansum' ) check_type(axis, 'axis', (int, list, tuple, type(None)), 'nansum') diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py index e16ac89953fb196bbe075012ac8f3317ff1eb7d6..2b6b4671efbffb7cec25320ad6ab8067f2f2fff9 100755 --- a/python/paddle/tensor/search.py +++ b/python/paddle/tensor/search.py @@ -574,8 +574,8 @@ def where(condition, x=None, y=None, name=None): Args: condition (Tensor): The condition to choose x or y. When True (nonzero), yield x, otherwise yield y. - x (Tensor|scalar, optional): A Tensor or scalar to choose when the condition is True with data type of float32, float64, int32 or int64. Either both or neither of x and y should be given. - y (Tensor|scalar, optional): A Tensor or scalar to choose when the condition is False with data type of float32, float64, int32 or int64. Either both or neither of x and y should be given. + x (Tensor|scalar, optional): A Tensor or scalar to choose when the condition is True with data type of float16, float32, float64, int32 or int64. Either both or neither of x and y should be given. + y (Tensor|scalar, optional): A Tensor or scalar to choose when the condition is False with data type of float16, float32, float64, int32 or int64. Either both or neither of x and y should be given. name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None. Returns: @@ -639,10 +639,10 @@ def where(condition, x=None, y=None, name=None): else: check_variable_and_dtype(condition, 'condition', ['bool'], 'where') check_variable_and_dtype( - x, 'x', ['float32', 'float64', 'int32', 'int64'], 'where' + x, 'x', ['float16', 'float32', 'float64', 'int32', 'int64'], 'where' ) check_variable_and_dtype( - y, 'y', ['float32', 'float64', 'int32', 'int64'], 'where' + y, 'y', ['float16', 'float32', 'float64', 'int32', 'int64'], 'where' ) helper = LayerHelper("where", **locals()) out = helper.create_variable_for_type_inference(dtype=x.dtype)