[0-size Tensor Job2 No.99] Add 0-size Tensor support for take_along_axis (PaddlePaddle#73832)

* fix

* add test of TakeAlongAxis

* fix test case

Author: zhengshengning

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()