support fp16 index sample

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

@@ -99,6 +99,28 @@ def config(self):
         self.index_type = "int64"
 
 
+class TestCase5(TestIndexSampleOp):
+    def config(self):
+        """
+        For float16 x type
+        """
+        self.x_shape = (10, 128)
+        self.x_type = "float16"
+        self.index_shape = (10, 64)
+        self.index_type = "int32"
+
+
+class TestCase6(TestIndexSampleOp):
+    def config(self):
+        """
+        For float16 x type
+        """
+        self.x_shape = (10, 128)
+        self.x_type = "float16"
+        self.index_shape = (10, 64)
+        self.index_type = "int64"
+
+
 class TestIndexSampleShape(unittest.TestCase):
     def test_shape(self):
         paddle.enable_static()

python/paddle/tensor/search.py

@@ -792,7 +792,7 @@ def index_sample(x, index):
             check_variable_and_dtype(
                 x,
                 'x',
-                ['float32', 'float64', 'int32', 'int64'],
+                ['float16', 'float32', 'float64', 'int32', 'int64'],
                 'paddle.tensor.search.index_sample',
             )
             check_variable_and_dtype(