PaddlePaddle
diff --git a/‎paddle/fluid/inference/tensorrt/plugin/gather_nd_op_plugin.cu
Copy file name to clipboard
+6-2Lines changed: 6 additions & 2 deletions b/‎paddle/fluid/inference/tensorrt/plugin/gather_nd_op_plugin.cu
Copy file name to clipboard
+6-2Lines changed: 6 additions & 2 deletions
diff --git a/‎paddle/fluid/prim/api/composite_backward/composite_backward_api.h
Copy file name to clipboard
+19-1Lines changed: 19 additions & 1 deletion b/‎paddle/fluid/prim/api/composite_backward/composite_backward_api.h
Copy file name to clipboard
+19-1Lines changed: 19 additions & 1 deletion
diff --git a/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+20-1Lines changed: 20 additions & 1 deletion b/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+20-1Lines changed: 20 additions & 1 deletion
diff --git a/‎paddle/phi/infermeta/binary.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/phi/infermeta/binary.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/phi/infermeta/ternary.cc
Copy file name to clipboard
+7-7Lines changed: 7 additions & 7 deletions b/‎paddle/phi/infermeta/ternary.cc
Copy file name to clipboard
+7-7Lines changed: 7 additions & 7 deletions
diff --git a/‎paddle/phi/kernels/cpu/gather_grad_kernel.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/phi/kernels/cpu/gather_grad_kernel.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/cpu/gather_kernel.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/phi/kernels/cpu/gather_kernel.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/cpu/index_add_impl.h
Copy file name to clipboard
+8-3Lines changed: 8 additions & 3 deletions b/‎paddle/phi/kernels/cpu/index_add_impl.h
Copy file name to clipboard
+8-3Lines changed: 8 additions & 3 deletions
diff --git a/‎paddle/phi/kernels/cpu/index_select_impl.h
Copy file name to clipboard
+11-3Lines changed: 11 additions & 3 deletions b/‎paddle/phi/kernels/cpu/index_select_impl.h
Copy file name to clipboard
+11-3Lines changed: 11 additions & 3 deletions
diff --git a/‎paddle/phi/kernels/funcs/gather.cu.h
Copy file name to clipboard
+26-13Lines changed: 26 additions & 13 deletions b/‎paddle/phi/kernels/funcs/gather.cu.h
Copy file name to clipboard
+26-13Lines changed: 26 additions & 13 deletions
@@ -48,13 +48,17 @@ __global__ void GatherNdCUDAKernel(const T* input,
     for (int32_t j = end_size - 1; j >= 0; --j) {
       auto index_value = indices[indices_i * end_size + j];
       PADDLE_ENFORCE(
-          index_value >= 0 && index_value < input_dims[j],
+          index_value >= -input_dims[j] && index_value < input_dims[j],
           "The index is out of bounds, "
           "please check whether the dimensions of index and "
           "input meet the requirements. It should "
-          "be less than [%d] and greater or equal to 0, but received [%d]",
+          "be less than [%d] and greater or equal to [%d], but received [%d]",
           input_dims[j],
+          -input_dims[j],
           index_value);
+      if (index_value < 0) {
+        index_value += input_dims[j];
+      }
       gather_i += (index_value * temp);
       temp *= input_dims[j];
     }
 
@@ -156,9 +156,13 @@ void gather_grad(const Tensor& x,
 
   // change axis to rank 0
   int axis_value = axis.to<int>();
+  int rank = x.dims().size();
+  if (axis_value < 0) {
+    axis_value += rank;
+  }
   tmp_perm.push_back(axis_value);
   // make other ranks
-  for (int i = 0; i < x.dims().size(); ++i) {
+  for (int i = 0; i < rank; ++i) {
     if (i != axis_value) {
       tmp_perm.push_back(i);
     }
@@ -1451,6 +1455,20 @@ void scatter_grad(const Tensor& index,
   if (updates_grad) {
     Scalar tmp_zero = 0;
     auto tmp_updates_grad = gather<T>(out_grad, index, tmp_zero);
+
+    // NOTE: len(index) can be smaller than len(updates) when updates is not a
+    // scalar
+    auto updates_dims = common::vectorize(updates.dims());
+    auto index_dims = common::vectorize(index.dims());
+    if (updates_dims.size() > 0 && updates_dims[0] > index_dims[0]) {
+      // Pad zeros to the end of tmp_updates_grad to make its shape the same as
+      // updates.
+      decltype(updates_dims) padding_dims = updates_dims;
+      padding_dims[0] = updates_dims[0] - index_dims[0];
+      auto padding_zeros = full<T>(padding_dims, 0, updates.dtype());
+      tmp_updates_grad =
+          concat<T>({tmp_updates_grad, std::move(padding_zeros)}, 0);
+    }
     set_output<T>(tmp_updates_grad, updates_grad);
   }
 }
 
@@ -475,6 +475,20 @@ void scatter_grad(const Tensor& index,
   if (updates_grad) {
     Scalar tmp_zero = 0;
     auto tmp_updates_grad = gather<T>(out_grad, index, tmp_zero);
+
+    // NOTE: len(index) can be smaller than len(updates) when updates is not a
+    // scalar
+    auto updates_dims = common::vectorize(updates.dims());
+    auto index_dims = common::vectorize(index.dims());
+    if (updates_dims.size() > 0 && updates_dims[0] > index_dims[0]) {
+      // Pad zeros to the end of tmp_updates_grad to make its shape the same as
+      // updates.
+      decltype(updates_dims) padding_dims = updates_dims;
+      padding_dims[0] = updates_dims[0] - index_dims[0];
+      auto padding_zeros = full<T>(padding_dims, 0, updates.dtype());
+      tmp_updates_grad =
+          concat<T>({tmp_updates_grad, std::move(padding_zeros)}, 0);
+    }
     set_output<T>(tmp_updates_grad, updates_grad);
   }
 }
@@ -1282,9 +1296,14 @@ void gather_grad(const Tensor& x,
 
   // change axis to rank 0
   int axis_value = axis.to<int>();
+  int rank = x.dims().size();
+  if (axis_value < 0) {
+    axis_value += rank;
+  }
+
   tmp_perm.push_back(axis_value);
   // make other ranks
-  for (int i = 0; i < x.dims().size(); ++i) {
+  for (int i = 0; i < rank; ++i) {
     if (i != axis_value) {
       tmp_perm.push_back(i);
     }
 
@@ -2024,7 +2024,7 @@ void GatherInferMeta(const MetaTensor& x,
 
   auto input_dim = x.dims();
   auto axis_v = axis.to<int>();
-  if (axis_v < 0) axis_v += input_dim.size();
+  if (axis_v < 0) axis_v += static_cast<int>(input_dim.size());
 
   PADDLE_ENFORCE_GE(
       axis_v,
 
@@ -1796,15 +1796,15 @@ void ScatterInferMeta(const MetaTensor& x,
             "Input(Updates)'s shape is %d.",
             ref_dims.size(),
             updates_dims.size()));
-    PADDLE_ENFORCE_EQ(
-        updates_dims[0],
+    PADDLE_ENFORCE_LE(
         index_dims[0],
+        updates_dims[0],
         common::errors::InvalidArgument(
-            "Input(Updates) and Input(Ids) should have same batch-size, but"
-            " received Input(Updates)'s batch-size is %d, Input(Ids)'s "
-            "batch-size is %d.",
-            updates_dims[0],
-            index_dims[0]));
+            "The first dimension size of Input(Index) shoud be no greater than "
+            "Input(Updates), but received first dimension size of Input(Index) "
+            "is %d, Input(Updates) is  %d.",
+            index_dims[0],
+            updates_dims[0]));
   } else {
     PADDLE_ENFORCE_EQ(
         (ref_dims.size() - 1 == updates_dims.size()),
 
@@ -31,6 +31,9 @@ void GatherGradKernel(const Context& dev_ctx,
                       DenseTensor* x_grad) {
   const auto& index_type = index.dtype();
   auto axis_v = axis.to<int>();
+  if (axis_v < 0) {
+    axis_v += static_cast<int>(x.dims().size());
+  }
 
   if (axis_v != 0) {
     if (index_type == phi::DataType::INT32) {
 
@@ -28,6 +28,11 @@ void GatherKernel(const Context& dev_ctx,
                   DenseTensor* out) {
   const auto& index_type = index.dtype();
   auto axis_v = axis.to<int>();
+  if (axis_v < 0) {
+    axis_v += static_cast<int>(x.dims().size());
+  }
+
+  // gather at non-zero axis
   if (axis_v != 0) {
     if (index_type == phi::DataType::INT32) {
       phi::funcs::GatherV2Function<T, int32_t>(
@@ -45,6 +50,7 @@ void GatherKernel(const Context& dev_ctx,
     return;
   }
 
+  // gather at axis 0
   if (index_type == phi::DataType::INT32) {
     phi::funcs::CPUGather<T, int>(dev_ctx, x, index, out);
   } else if (index_type == phi::DataType::INT64) {
 
@@ -56,20 +56,22 @@ void IndexAddInner(const Context& ctx,
   for (int i = 0; i < index_size; i++) {
     PADDLE_ENFORCE_GE(
         index_data[i],
-        0,
+        -input_dim[axis],
         common::errors::InvalidArgument(
             "Variable value (index) of OP(index_add) "
-            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "expected >= %ld and < %ld, but got %ld. Please check input "
             "value.",
+            -input_dim[axis],
             input_dim[axis],
             index_data[i]));
     PADDLE_ENFORCE_LT(
         index_data[i],
         input_dim[axis],
         common::errors::InvalidArgument(
             "Variable value (index) of OP(index_add) "
-            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "expected >= %ld and < %ld, but got %ld. Please check input "
             "value.",
+            -input_dim[axis],
             input_dim[axis],
             index_data[i]));
   }
@@ -88,6 +90,9 @@ void IndexAddInner(const Context& ctx,
   auto& place = *ctx.eigen_device();
   for (auto j = 0; j < index_size; j++) {
     IndexT index_value = index_data[j];
+    if (index_value < 0) {
+      index_value += input_dim[axis];
+    }
     auto output_t = output_tensor.chip(index_value, 1);
     output_t.device(place) = output_t + add_value_tensor.chip(j, 1);
   }
 
@@ -85,20 +85,22 @@ void IndexSelectInner(const Context& ctx,
   for (int i = 0; i < index_size; i++) {
     PADDLE_ENFORCE_GE(
         index_data[i],
-        0,
+        -input_dim[dim],
         common::errors::InvalidArgument(
             "Variable value (index) of OP(index_select) "
-            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "expected >= %ld and < %ld, but got %ld. Please check input "
             "value.",
+            -input_dim[dim],
             input_dim[dim],
             index_data[i]));
     PADDLE_ENFORCE_LT(
         index_data[i],
         input_dim[dim],
         common::errors::InvalidArgument(
             "Variable value (index) of OP(index_select) "
-            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "expected >= %ld and < %ld, but got %ld. Please check input "
             "value.",
+            -input_dim[dim],
             input_dim[dim],
             index_data[i]));
   }
@@ -116,6 +118,9 @@ void IndexSelectInner(const Context& ctx,
 
   for (auto j = 0; j < index_size; j++) {
     IndexT index_value = index_data[j];
+    if (index_value < 0) {
+      index_value += input_dim[dim];
+    }
     auto output_t = output_tensor.chip(j, 1);
     output_t.device(place) = input_tensor.chip(index_value, 1);
   }
@@ -167,6 +172,9 @@ void IndexSelectGradInner(const Context& ctx,
 
     for (auto j = 0; j < index_size; j++) {
       IndexT index_value = index_data[j];
+      if (index_value < 0) {
+        index_value += input_dim[dim];
+      }
       auto src = input_data + input_start_offset + j * slice_size;
       auto p_out = p_output + output_start_offset + index_value * slice_size;
       auto dst = out_data + output_start_offset + index_value * slice_size;
 
@@ -32,11 +32,14 @@ __global__ void GatherCUDAKernel(const T* params,
                                  const IndexT* indices,
                                  T* output,
                                  size_t index_size,
-                                 size_t slice_size) {
+                                 size_t slice_size,
+                                 int64_t index_dim_size) {
   CUDA_KERNEL_LOOP_TYPE(i, index_size * slice_size, int64_t) {
     int64_t indices_i = i / slice_size;
     int64_t slice_i = i - indices_i * slice_size;  // offset inside the slice
-    IndexT gather_i = indices[indices_i];
+    IndexT gather_i =
+        (indices[indices_i] < 0 ? (indices[indices_i] + index_dim_size)
+                                : indices[indices_i]);
     int64_t params_i = gather_i * slice_size + slice_i;
     *(output + i) = *(params + params_i);
   }
@@ -57,14 +60,19 @@ __global__ void GatherNdCUDAKernel(const T* input,
     int64_t temp = slice_size;
     for (int64_t j = end_size - 1; j >= 0; --j) {
       auto index_value = indices[indices_i * end_size + j];
-      PADDLE_ENFORCE(index_value >= 0 && index_value < input_dims[j],
-                     "The index is out of bounds, "
-                     "please check whether the dimensions of index and "
-                     "input meet the requirements. It should "
-                     "be less than [%ld] and greater than or equal to 0, but "
-                     "received [%ld]",
-                     input_dims[j],
-                     index_value);
+      PADDLE_ENFORCE(
+          index_value >= -input_dims[j] && index_value < input_dims[j],
+          "The index is out of bounds, "
+          "please check whether the dimensions of index and "
+          "input meet the requirements. It should "
+          "be less than [%ld] and greater than or equal to [%ld], but "
+          "received [%ld]",
+          input_dims[j],
+          -input_dims[j],
+          index_value);
+      if (index_value < 0) {
+        index_value += input_dims[j];
+      }
       gather_i += (index_value * temp);
       temp *= input_dims[j];
     }
@@ -112,7 +120,7 @@ void GPUGather(const phi::GPUContext& ctx,
   phi::backends::gpu::LimitGridDim(ctx, &grid);
 
   GatherCUDAKernel<T, IndexT><<<grid, block, 0, ctx.stream()>>>(
-      p_src, p_index, p_output, index_size, slice_size);
+      p_src, p_index, p_output, index_size, slice_size, src_dims[0]);
 }
 
 template <typename T, typename IndexT = int>
@@ -180,13 +188,18 @@ __global__ void GatherGPUKernel(const T* input,
     U index_val = index[index_dim_index];
 
     PADDLE_ENFORCE(
-        index_val >= 0 && index_val < input_index_dim_size,
+        index_val >= -input_index_dim_size && index_val < input_index_dim_size,
         "The index is out of bounds, "
         "please check whether the dimensions of index and "
         "input meet the requirements. It should "
-        "be less than [%ld] and greater than or equal to 0, but received [%ld]",
+        "be less than [%ld] and greater than or equal to [%ld], but "
+        "received [%ld]",
         input_index_dim_size,
+        -input_index_dim_size,
         index_val);
+    if (index_val < 0) {
+      index_val += input_index_dim_size;
+    }
 
     int64_t out_dim_index = next_idx - outer_dim_size * index_dim_index;
     int64_t input_index =