From 555c60e214a872014da29cdb790ca42398879a1d Mon Sep 17 00:00:00 2001
From: ooooo <3164076421@qq.com>
Date: Fri, 13 Jun 2025 20:42:33 +0800
Subject: [PATCH] add tests

---
 paddle/phi/kernels/funcs/gather.h  |  6 +++--
 test/legacy_test/test_gather_op.py | 40 ++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/paddle/phi/kernels/funcs/gather.h b/paddle/phi/kernels/funcs/gather.h
index e1609cac7f0546..cc45f18361a75a 100644
--- a/paddle/phi/kernels/funcs/gather.h
+++ b/paddle/phi/kernels/funcs/gather.h
@@ -265,12 +265,14 @@ void GatherV2GradFunction(const phi::CPUContext& ctx,
   for (int64_t i = 0; i < inner_dim_size; i++) {
     for (int64_t j = 0; j < input_index_dim_size; j++) {
       const int64_t index_data_j =
-          (index_data[j] < 0 ? index_data[j] + input_index_dim_size
+          (index_data[j] < 0 ? index_data[j] + out_index_dim_size
                              : index_data[j]);
       for (int64_t k = 0; k < outer_dim_size; k++) {
         int64_t index = k + index_data_j * outer_dim_size +
                         i * outer_dim_size * out_index_dim_size;
-        out_data[index] += input_data[j * outer_dim_size + k];
+        out_data[index] +=
+            input_data[i * input_index_dim_size * outer_dim_size +
+                       j * outer_dim_size + k];
       }
     }
   }
diff --git a/test/legacy_test/test_gather_op.py b/test/legacy_test/test_gather_op.py
index d8227134d6b5d2..cdb5c9e8488113 100644
--- a/test/legacy_test/test_gather_op.py
+++ b/test/legacy_test/test_gather_op.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy as np
@@ -908,6 +909,45 @@ def test_pir_out_type(self):
             self.assertTrue(out.dtype == core.DataType.INT64)
 
 
+class TestGatherBackward(unittest.TestCase):
+    def setUp(self):
+        self.shape = [10, 20]
+        self.dtype = 'float32'
+        self.index = (1, 3, 5)
+        self.index_dtype = 'int64'
+        self.places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            self.places.append(paddle.CPUPlace())
+        if paddle.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+
+    def test_gather_backward(self):
+        if len(self.places) != 2:
+            return
+        res_list = []
+        x_np = np.random.random(self.shape).astype(self.dtype)
+        index_np = np.array(self.index, dtype=self.index_dtype)
+        grad_out_np = np.random.random(self.shape).astype(self.dtype)
+        for place in self.places:
+            with base.dygraph.guard(place):
+                x = paddle.to_tensor(x_np, dtype=self.dtype)
+                x.stop_gradient = False
+                index = paddle.to_tensor(index_np, dtype=self.index_dtype)
+                out = paddle.gather(x, index, -1)
+                grad_out = paddle.to_tensor(grad_out_np, dtype=self.dtype)
+                (re,) = paddle.grad(
+                    outputs=out,
+                    inputs=x,
+                    grad_outputs=grad_out,
+                )
+                res_list.append(re.numpy())
+        np.testing.assert_allclose(res_list[0], res_list[1])
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()