[XPU] support take_along_axis_grad in XPU (#71779)

* xpu: support take_along_axis_grad in XPU * fix test cases for take_along_axis

Author: zhouquan32

paddle/phi/backends/xpu/xpu3_op_list.cc

@@ -1497,6 +1497,10 @@ XPUOpMap& get_kl3_ops() {
  XPUKernelSet({phi::DataType::FLOAT32,
  phi::DataType::FLOAT16,
  phi::DataType::BFLOAT16})},
+ {"take_along_axis_grad",
+ XPUKernelSet({phi::DataType::FLOAT32,
+ phi::DataType::FLOAT16,
+ phi::DataType::BFLOAT16})},
  {"tanh_grad",
  XPUKernelSet({phi::DataType::FLOAT32,
  phi::DataType::FLOAT16,

paddle/phi/kernels/xpu/take_along_axis_grad_kernel.cc

@@ -0,0 +1,92 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/take_along_axis_grad_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void TakeAlongAxisGradKernel(const Context& dev_ctx,
+ const DenseTensor& x,
+ const DenseTensor& index,
+ const DenseTensor& out_grad,
+ int axis,
+ DenseTensor* x_grad) {
+ using XPUType = typename XPUTypeTrait<T>::Type;
+ dev_ctx.template Alloc<T>(x_grad);
+
+ const auto& index_dtype = index.dtype();
+ bool index_dtype_match =
+ index_dtype == DataType::INT32 || index_dtype == DataType::INT64;
+ PADDLE_ENFORCE_EQ(index_dtype_match,
+ true,
+ errors::InvalidArgument(
+ "Input(Index) holds the wrong type, it holds %s, but "
+ "desires to be %s or %s",
+ DataTypeToString(index_dtype),
+ DataTypeToString(DataType::INT32),
+ DataTypeToString(DataType::INT64)));
+
+ int r = xpu::constant(dev_ctx.x_context(),
+ reinterpret_cast<XPUType*>(x_grad->data<T>()),
+ x_grad->numel(),
+ XPUType(0));
+ PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
+
+ auto x_shape = common::vectorize<int64_t>(x.dims());
+ auto out_grad_shape = common::vectorize<int64_t>(out_grad.dims());
+ auto index_shape = common::vectorize<int64_t>(index.dims());
+
+ if (index_dtype == DataType::INT32) {
+ r = xpu::paddle_put_along_axis<XPUType, int>(
+ dev_ctx.x_context(),
+ reinterpret_cast<const XPUType*>(x_grad->data<T>()),
+ reinterpret_cast<const XPUType*>(out_grad.data<T>()),
+ reinterpret_cast<const int*>(index.data<int>()),
+ reinterpret_cast<XPUType*>(x_grad->data<T>()),
+ x_shape,
+ out_grad_shape,
+ index_shape,
+ axis,
+ 1,
+ false);
+ PADDLE_ENFORCE_XDNN_SUCCESS(r, "paddle_put_along_axis");
+ } else {
+ r = xpu::paddle_put_along_axis<XPUType, int64_t>(
+ dev_ctx.x_context(),
+ reinterpret_cast<const XPUType*>(x_grad->data<T>()),
+ reinterpret_cast<const XPUType*>(out_grad.data<T>()),
+ reinterpret_cast<const int64_t*>(index.data<int64_t>()),
+ reinterpret_cast<XPUType*>(x_grad->data<T>()),
+ x_shape,
+ out_grad_shape,
+ index_shape,
+ axis,
+ 1,
+ false);
+ PADDLE_ENFORCE_XDNN_SUCCESS(r, "paddle_put_along_axis");
+ }
+}
+} // namespace phi
+
+PD_REGISTER_KERNEL(take_along_axis_grad,
+ XPU,
+ ALL_LAYOUT,
+ phi::TakeAlongAxisGradKernel,
+ float,
+ phi::dtype::float16,
+ phi::dtype::bfloat16) {}

test/xpu/test_take_along_axis_op_xpu.py

@@ -61,7 +61,6 @@ def setUp(self):
  self.outputs = {'Result': self.target}
 
  def init_config(self):
- self.in_type = np.float32
  self.x_shape = (1, 4, 10)
  self.index_type = np.int32
  self.index = np.array([[[0, 1, 3, 5, 6]]]).astype(self.index_type)
@@ -77,39 +76,38 @@ def test_check_grad(self):
 
  class TestCase1(TestXPUTakeAlongAxisOp):
  def init_config(self):
- self.in_type = np.float32
  self.x_shape = (1, 10, 100)
  self.index_type = np.int32
  self.index = np.array([[[0, 1, 3, 5, 13]]]).astype(self.index_type)
  self.axis = 2
 
  class TestCase2(TestXPUTakeAlongAxisOp):
  def init_config(self):
- self.in_type = np.float32
  self.x_shape = (1, 10, 100)
  self.index_type = np.int64
  self.index = np.array([[[0, 1, 3, 5, 13]]]).astype(self.index_type)
  self.axis = 2
 
  class TestCase3(TestXPUTakeAlongAxisOp):
  def init_config(self):
- self.in_type = np.float16
  self.x_shape = (1, 10, 100)
  self.index_type = np.int32
- self.index = np.array([[[0, 1, 3, 5, 13]]]).astype(self.index_type)
- self.axis = 2
+ self.index = np.array([[[0], [1], [3], [5]]]).astype(
+ self.index_type
+ )
+ self.axis = 1
 
  class TestCase4(TestXPUTakeAlongAxisOp):
  def init_config(self):
- self.in_type = np.float16
  self.x_shape = (1, 10, 100)
  self.index_type = np.int64
- self.index = np.array([[[0, 1, 3, 5, 13]]]).astype(self.index_type)
- self.axis = 2
+ self.index = np.array([[[0], [1], [3], [5]]]).astype(
+ self.index_type
+ )
+ self.axis = 1
 
  class TestCase5(TestXPUTakeAlongAxisOp):
  def init_config(self):
- self.in_type = np.float32
  self.x_shape = (1, 10, 100)
  self.index_type = np.int32
  self.index = np.array([[[0], [1], [3], [5], [8]]]).astype(
@@ -119,8 +117,7 @@ def init_config(self):
 
  class TestCase6(TestXPUTakeAlongAxisOp):
  def init_config(self):
- self.in_type = np.uint16
- self.x_shape = (1, 10, 100)
+ self.x_shape = (2, 30, 50)
  self.index_type = np.int64
  self.index = np.array([[[0, 1, 3, 5, 13]]]).astype(self.index_type)
  self.axis = 2