[AMP OP&Test]add fp16 and bf16 OpTest for index_select (#51159)

* add fp16/bf16

* add grad bf16

* test name

paddle/phi/kernels/cpu/index_select_kernel.cc

@@ -58,5 +58,6 @@ PD_REGISTER_KERNEL(index_select,
                    phi::IndexSelectKernel,
                    float,
                    double,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}

paddle/phi/kernels/gpu/index_select_grad_kernel.cu

@@ -132,5 +132,6 @@ PD_REGISTER_KERNEL(index_select_grad,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}

paddle/phi/kernels/gpu/index_select_kernel.cu

@@ -84,5 +84,6 @@ PD_REGISTER_KERNEL(index_select,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}

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

@@ -15,7 +15,7 @@
 import unittest
 
 import numpy as np
-from op_test import OpTest
+from op_test import OpTest, convert_float_to_uint16
 
 import paddle
 import paddle.fluid as fluid
@@ -81,6 +81,54 @@ def init_dtype_type(self):
         self.index_size = 10
 
 
+class TestIndexSelectFP16OP(TestIndexSelectOp):
+    def init_dtype_type(self):
+        self.dim = 1
+        self.x_type = np.float16
+        self.index_type = np.int64
+        self.x_shape = (100, 4, 5)
+        self.index_size = 100
+
+
+class TestIndexSelectBF16Op(OpTest):
+    def setUp(self):
+        self.python_api = paddle.index_select
+        self.op_type = "index_select"
+        self.init_dtype_type()
+        index_np = np.random.randint(
+            low=0, high=self.x_shape[self.dim], size=self.index_size
+        )
+        x_np = np.random.random(self.x_shape).astype(np.float32)
+        self.inputs = {'X': convert_float_to_uint16(x_np), 'Index': index_np}
+        self.attrs = {'dim': self.dim}
+        outer_loop = np.prod(self.x_shape[: self.dim])
+        x_reshape = [outer_loop] + list(self.x_shape[self.dim :])
+        x_np_reshape = np.reshape(x_np, tuple(x_reshape))
+        out_list = []
+        for i in range(outer_loop):
+            for j in range(self.index_size):
+                out_list.append(x_np_reshape[i, index_np[j]])
+        self.out_shape = list(self.x_shape)
+        self.out_shape[self.dim] = self.index_size
+        self.out_shape = tuple(self.out_shape)
+
+        out = np.reshape(out_list, self.out_shape)
+        self.outputs = {'Out': convert_float_to_uint16(out)}
+
+    def init_dtype_type(self):
+        self.dim = 1
+        self.x_type = np.uint16
+        self.index_type = np.int64
+        self.x_shape = (100, 4, 5)
+        self.index_size = 100
+
+    def test_check_output(self):
+        self.check_output(check_eager=True)
+
+    def test_check_grad_normal(self):
+        self.check_grad(['X'], 'Out', check_eager=True)
+
+
 class TestIndexSelectAPI(unittest.TestCase):
     def input_data(self):
         self.data_x = np.array(

python/paddle/tensor/search.py

@@ -350,7 +350,7 @@ def index_select(x, index, axis=0, name=None):
         check_variable_and_dtype(
             x,
             'x',
-            ['float16', 'float32', 'float64', 'int32', 'int64'],
+            ['bfloat16', 'float16', 'float32', 'float64', 'int32', 'int64'],
             'paddle.tensor.search.index_select',
         )
         check_variable_and_dtype(