fix fp16 dtype checking for argsort op (#50939)

python/paddle/fluid/tests/unittests/test_argsort_op.py

@@ -499,5 +499,18 @@ class TestArgsortWithInputNaN(unittest.TestCase):
         paddle.enable_static()
 
 
+class TestArgsortOpFp16(unittest.TestCase):
+    def test_fp16(self):
+        x_np = np.random.random((2, 8)).astype('float16')
+        with paddle.static.program_guard(paddle.static.Program()):
+            x = paddle.static.data(shape=[2, 8], name='x', dtype='float16')
+            out = paddle.argsort(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])
+
+
 if __name__ == "__main__":
     unittest.main()

python/paddle/tensor/search.py

@@ -39,7 +39,7 @@ def argsort(x, axis=-1, descending=False, name=None):
     Sorts the input along the given axis, and returns the corresponding index tensor for the sorted output values. The default sort algorithm is ascending, if you want the sort algorithm to be descending, you must set the :attr:`descending` as True.
 
     Args:
-        x(Tensor): An input N-D Tensor with type float32, float64, int16,
+        x(Tensor): An input N-D Tensor with type float16, float32, float64, int16,
             int32, int64, uint8.
         axis(int, optional): Axis to compute indices along. The effective range
             is [-R, R), where R is Rank(x). when axis<0, it works the same way
@@ -101,7 +101,15 @@ def argsort(x, axis=-1, descending=False, name=None):
         check_variable_and_dtype(
             x,
             'x',
-            ['float32', 'float64', 'int16', 'int32', 'int64', 'uint8'],
+            [
+                'float16',
+                'float32',
+                'float64',
+                'int16',
+                'int32',
+                'int64',
+                'uint8',
+            ],
             'argsort',
         )