[0-size Tensor Job2 No.99] Add 0-size Tensor support for take_along_axis

paddle/phi/kernels/cpu/take_along_axis_grad_kernel.cc

@@ -34,6 +34,10 @@ void TakeAlongAxisGradKernel(const Context& dev_ctx,
  x_grad->Resize(x.dims());
  dev_ctx.template Alloc<T>(x_grad);
 
+ if (x_grad->numel() == 0) {
+ return;
+ }
+
  // Set to zero tensor.
  phi::funcs::SetConstant<Context, T> functor;
  functor(dev_ctx, x_grad, static_cast<T>(0));

paddle/phi/kernels/cpu/take_along_axis_kernel.cc

@@ -42,6 +42,10 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
  out->Resize(index.dims());
  dev_ctx.template Alloc<T>(out);
 
+ if (out->numel() == 0) {
+ return;
+ }
+
  const auto& index_type = index.dtype();
  if (index_type == DataType::INT32) {
  phi::funcs::cpu_gather_kernel<T, int32_t>(

paddle/phi/kernels/gpu/take_along_axis_grad_kernel.cu

@@ -35,6 +35,10 @@ void TakeAlongAxisGradKernel(const Context& dev_ctx,
  x_grad->Resize(x.dims());
  dev_ctx.template Alloc<T>(x_grad);
 
+ if (x_grad->numel() == 0) {
+ return;
+ }
+
  // Set to zero tensor.
  phi::funcs::SetConstant<Context, T> functor;
  functor(dev_ctx, x_grad, static_cast<T>(0));

paddle/phi/kernels/gpu/take_along_axis_kernel.cu

@@ -42,6 +42,10 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
  out->Resize(index.dims());
  dev_ctx.template Alloc<T>(out);
 
+ if (out->numel() == 0) {
+ return;
+ }
+
  const auto& index_type = index.dtype();
  if (index_type == DataType::INT32) {
  phi::funcs::gpu_gather_kernel<T, int32_t>(

paddle/phi/kernels/xpu/take_along_axis_grad_kernel.cc

@@ -29,6 +29,10 @@ void TakeAlongAxisGradKernel(const Context& dev_ctx,
  using XPUType = typename XPUTypeTrait<T>::Type;
  dev_ctx.template Alloc<T>(x_grad);
 
+ if (x_grad->numel() == 0) {
+ return;
+ }
+
  const auto& index_dtype = index.dtype();
  bool index_dtype_match =
  index_dtype == DataType::INT32 || index_dtype == DataType::INT64;

paddle/phi/kernels/xpu/take_along_axis_kernel.cc

@@ -42,6 +42,10 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
  out->Resize(index.dims());
  dev_ctx.template Alloc<T>(out);
 
+ if (out->numel() == 0) {
+ return;
+ }
+
  if (x.numel() == 0 || index.numel() == 0) return;
 
  const auto& index_dtype = index.dtype();

test/legacy_test/test_take_along_axis_op.py

@@ -428,6 +428,41 @@ def test_static_shape_take_along_axis(self):
  _ = static_f(x, ind, axis=0, broadcast=False)
 
 
+class TestTakeAlongAxis_ZeroSize(OpTest):
+ def setUp(self):
+ self.python_api = paddle.take_along_axis
+ self.op_type = "take_along_axis"
+ self.dtype = "float64"
+ self.check_pir = True
+
+ x = np.zeros((2, 0, 5)).astype(self.dtype)
+ indices = np.zeros((2, 3, 5)).astype("int64")
+
+ self.inputs = {'Input': x, 'Index': indices}
+ self.attrs = {'Axis': 1}
+
+ output = np.zeros((2, 3, 5)).astype(self.dtype)
+ self.outputs = {'Result': output}
+
+ def test_check_output(self):
+ self.check_output_with_place(
+ paddle.CPUPlace(), check_pir=self.check_pir
+ )
+ if core.is_compiled_with_cuda():
+ self.check_output_with_place(
+ core.CUDAPlace(0), check_pir=self.check_pir
+ )
+
+ def test_check_grad(self):
+ self.check_grad_with_place(
+ paddle.CPUPlace(), ['Input'], 'Result', check_pir=self.check_pir
+ )
+ if core.is_compiled_with_cuda():
+ self.check_grad_with_place(
+ core.CUDAPlace(0), ['Input'], 'Result', check_pir=self.check_pir
+ )
+
+
 if __name__ == "__main__":
  paddle.enable_static()
  unittest.main()