[fp16] suppot fp16 input in nansum (#50847)

* add float16 in python/paddle/math

* add unittest for float16

* add float16 support in python.paddle.tensor.search.where

* remove fp16 error cases

* Add NotImplementedError unittest

* fix codestyle

* fluid to paddle.static; add cases with GPU

* Add float16 in English docs

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()
+        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)
 
-        # input dtype error
-        def run1():
-            input = fluid.data(name='input', dtype='float16', shape=[2, 3])
-            output = paddle.nansum(input)
-
-        self.assertRaises(TypeError, run1)
+                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],
+                )
 
-        # axis type error
-        def run2():
-            input = fluid.data(name='input', dtype='float16', shape=[2, 3])
-            output = paddle.nansum(input, axis=1.2)
+                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.assertRaises(TypeError, run2)
+                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(

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')
 

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)