[0-size Tensor No.271] Add 0-size Tensor support for paddle.take_along_axis API. (#73736)

Author: straigrand

paddle/phi/kernels/cpu/take_along_axis_kernel.cc

@@ -18,6 +18,7 @@
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/gather_scatter_functor.h"
 
 namespace phi {
@@ -28,6 +29,16 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
                          const DenseTensor& index,
                          int axis,
                          DenseTensor* out) {
+  if (index.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, common::vectorize(out->dims()), static_cast<T>(0), out);
+    return;
+  }
+
   out->Resize(index.dims());
   dev_ctx.template Alloc<T>(out);
 

paddle/phi/kernels/gpu/take_along_axis_kernel.cu

@@ -18,6 +18,7 @@
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/utils/data_type.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/gather_scatter_functor.h"
 
 namespace phi {
@@ -28,6 +29,16 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
                          const DenseTensor& index,
                          int axis,
                          DenseTensor* out) {
+  if (index.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, common::vectorize(out->dims()), static_cast<T>(0), out);
+    return;
+  }
+
   out->Resize(index.dims());
   dev_ctx.template Alloc<T>(out);
 

paddle/phi/kernels/xpu/take_along_axis_kernel.cc

@@ -19,6 +19,7 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 
@@ -28,6 +29,16 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
                          const DenseTensor& index,
                          int axis,
                          DenseTensor* out) {
+  if (index.numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, common::vectorize(out->dims()), static_cast<T>(0), out);
+    return;
+  }
+
   out->Resize(index.dims());
   dev_ctx.template Alloc<T>(out);
 

test/legacy_test/test_take_along_axis_op.py

@@ -25,6 +25,58 @@
 paddle.enable_static()
 
 
+class TestTakeAlongAxis0Size(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(check_pir=self.check_pir)
+
+    def test_check_grad(self):
+        self.check_grad(['Input'], 'Result', check_pir=self.check_pir)
+
+
+class TestTakeAlongAxis0Size2(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.random.rand(2, 3, 5).astype(self.dtype)
+        indices = np.zeros((2, 0, 5)).astype("int64")
+
+        self.inputs = {'Input': x, 'Index': indices}
+        self.attrs = {'Axis': 1}
+
+        output = np.zeros((2, 0, 5)).astype(self.dtype)
+        self.outputs = {'Result': output}
+
+    def test_check_output(self):
+        self.check_output(check_pir=self.check_pir)
+
+    def test_check_grad(self):
+        self.grad = np.zeros_like(self.outputs['Result']).astype(self.dtype)
+        self.check_grad(
+            ['Input'],
+            'Result',
+            user_defined_grads=[self.grad],
+            check_pir=self.check_pir,
+        )
+
+
 class TestTakeAlongAxisOp(OpTest):
     def setUp(self):
         self.init_data()