Comp index select (#51215)

python/paddle/fluid/tests/unittests/CMakeLists.txt

@@ -1199,6 +1199,7 @@ set(TEST_CINN_OPS
     test_stack_op
     test_activation_op
     test_full_like_op
+    test_index_select_op
     test_fill_any_like_op
     test_concat_op
     test_top_k_v2_op

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

@@ -28,7 +28,9 @@ class TestIndexSelectOp(OpTest):
     def setUp(self):
         self.python_api = paddle.index_select
         self.op_type = "index_select"
+        self.prim_op_type = "comp"
         self.init_dtype_type()
+
         index_np = np.random.randint(
             low=0, high=self.x_shape[self.dim], size=self.index_size
         )
@@ -57,10 +59,10 @@ class TestIndexSelectOp(OpTest):
         self.index_size = 100
 
     def test_check_output(self):
-        self.check_output(check_eager=True)
+        self.check_output(check_eager=True, check_prim=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X'], 'Out', check_eager=True)
+        self.check_grad(['X'], 'Out', check_eager=True, check_prim=True)
 
 
 class TestIndexSelectOpCase2(TestIndexSelectOp):
@@ -92,6 +94,7 @@ class TestIndexSelectFP16OP(TestIndexSelectOp):
         self.index_size = 100
 
 
+# no scatter op (the backward op of index_select/gather) for bf16
 class TestIndexSelectBF16Op(OpTest):
     def setUp(self):
         self.python_api = paddle.index_select

python/paddle/incubate/autograd/composite_rules.py

@@ -393,6 +393,15 @@ def hard_swish_composite(x):
     return res
 
 
+@REGISTER_COMPOSITE('index_select')
+def index_select_composite(x, index, axis):
+    """define composite rule of op index_select."""
+    if axis < 0:
+        axis = len(x.shape) + axis
+    res = gather(x, index, axis=axis)
+    return res
+
+
 @REGISTER_COMPOSITE('sigmoid')
 def sigmoid_composite(x):
     """

python/paddle/incubate/autograd/primitives.py

@@ -34,6 +34,7 @@ from paddle.tensor import erfinv  # noqa: F401
 from paddle.tensor import exp  # noqa: F401
 from paddle.tensor import expm1  # noqa: F401
 from paddle.tensor import full  # noqa: F401
+from paddle.tensor import gather  # noqa: F401
 from paddle.tensor import greater_equal  # noqa: F401
 from paddle.tensor import lgamma  # noqa: F401
 from paddle.tensor import log  # noqa: F401
@@ -124,6 +125,7 @@ others = [
     'cast',
     'fill_constant',
     'reshape',
+    'gather'
     'full',
     'tile',
     'concat',