[0-size Tensor No.270] Add 0-size Tensor support for take (#73186)

co63oc · web-flow · commit 200c28d94ba8 · 2025-06-11T14:22:48.000+08:00
* Fix

* Fix
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2434,13 +2434,6 @@ void IndexSelectInferMeta(const MetaTensor& x,
     dim += input_dim.size();
   }
 
-  if (input_dim[dim] != 0) {
-    PADDLE_ENFORCE_EQ(index_dim[0] != 0,
-                      true,
-                      common::errors::InvalidArgument(
-                          "The length of Input(Index) can't be 0."));
-  }
-
   auto output_dim = common::vectorize(input_dim);
 
   output_dim[dim] = index_dim[0];
diff --git a/paddle/phi/kernels/cpu/index_select_grad_kernel.cc b/paddle/phi/kernels/cpu/index_select_grad_kernel.cc
@@ -17,6 +17,7 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/utils/data_type.h"
 #include "paddle/phi/kernels/cpu/index_select_impl.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 
@@ -27,6 +28,12 @@ void IndexSelectGradKernel(const Context& dev_ctx,
                            const DenseTensor& out_grad,
                            int dim,
                            DenseTensor* x_grad) {
+  // x [3, 4], index [5, 0], out [5, 0]
+  if (out_grad.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(x.dims())), 0, x_grad);
+    return;
+  }
   if (dim < 0) {
     dim += out_grad.dims().size();
   }
diff --git a/paddle/phi/kernels/cpu/index_select_kernel.cc b/paddle/phi/kernels/cpu/index_select_kernel.cc
@@ -26,6 +26,10 @@ void IndexSelectKernel(const Context& dev_ctx,
                        const DenseTensor& index,
                        int dim,
                        DenseTensor* output) {
+  if (output && output->numel() == 0) {
+    dev_ctx.template Alloc<T>(output);
+    return;
+  }
   auto inputs = x;
   if (dim < 0) {
     dim += inputs.dims().size();
diff --git a/paddle/phi/kernels/gpu/index_select_grad_kernel.cu b/paddle/phi/kernels/gpu/index_select_grad_kernel.cu
@@ -21,6 +21,7 @@
 #include "paddle/phi/backends/gpu/gpu_primitives.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/math_function.h"
 
 COMMON_DECLARE_bool(cudnn_deterministic);
@@ -58,6 +59,12 @@ void IndexSelectGradKernel(const Context& dev_ctx,
                            const DenseTensor& out_grad,
                            int dim,
                            DenseTensor* x_grad) {
+  if (out_grad.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(x.dims())), 0, x_grad);
+    return;
+  }
+
   auto* output_grad_data = out_grad.data<T>();
   auto* in_grad_data = dev_ctx.template Alloc<T>(x_grad);
 
diff --git a/paddle/phi/kernels/gpu/index_select_kernel.cu b/paddle/phi/kernels/gpu/index_select_kernel.cu
@@ -31,6 +31,10 @@ void IndexSelectKernel(const Context& dev_ctx,
                        const DenseTensor& index,
                        int dim,
                        DenseTensor* output) {
+  if (output && output->numel() == 0) {
+    dev_ctx.template Alloc<T>(output);
+    return;
+  }
   auto input_dim = x.dims();
   auto output_dim = output->dims();
   dim = dim >= 0 ? dim : dim + input_dim.size();
@@ -56,9 +60,6 @@ void IndexSelectKernel(const Context& dev_ctx,
   T* out_data = dev_ctx.template Alloc<T>(output);
 
   int64_t numel = output->numel();
-  if (numel == 0) {
-    return;
-  }
   auto stream = dev_ctx.stream();
 
   unsigned int block_dim = PADDLE_CUDA_NUM_THREADS;
diff --git a/paddle/phi/kernels/xpu/index_select_grad_kernel.cc b/paddle/phi/kernels/xpu/index_select_grad_kernel.cc
@@ -17,7 +17,7 @@
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/utils/data_type.h"
-
+#include "paddle/phi/kernels/full_kernel.h"
 namespace phi {
 
 template <typename T, typename Context>
@@ -28,6 +28,11 @@ void IndexSelectGradKernel(const Context& dev_ctx,
                            int dim,
                            DenseTensor* x_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
+  if (out_grad.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(x.dims())), 0, x_grad);
+    return;
+  }
   if (dim < 0) {
     dim += out_grad.dims().size();
   }
diff --git a/paddle/phi/kernels/xpu/index_select_kernel.cc b/paddle/phi/kernels/xpu/index_select_kernel.cc
@@ -27,6 +27,10 @@ void IndexSelectKernel(const Context& dev_ctx,
                        const DenseTensor& index,
                        int dim,
                        DenseTensor* output) {
+  if (output && output->numel() == 0) {
+    dev_ctx.template Alloc<T>(output);
+    return;
+  }
   auto input_dim = x.dims();
   dim = dim >= 0 ? dim : dim + input_dim.size();
   const auto& index_type = index.dtype();
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -7162,7 +7162,8 @@ def take(
         )
     elif mode == 'clip':
         # 'clip' mode disables indexing with negative numbers.
-        index_1d = clip(index_1d, 0, max_index - 1)
+        if max_index > 0:  # If max_index is 0, input_1d is 0-size.
+            index_1d = clip(index_1d, 0, max_index - 1)
 
     out = input_1d.index_select(index_1d).reshape(index.shape)
 
diff --git a/test/legacy_test/test_take.py b/test/legacy_test/test_take.py
@@ -269,5 +269,48 @@ def set_input(self):
         )  # Both ends of the index are out of bounds
 
 
+class TestTakeAPI_ZeroSize(unittest.TestCase):
+    def set_mode(self):
+        self.mode = 'raise'
+
+    def set_dtype(self):
+        self.input_dtype = 'float64'
+        self.index_dtype = 'int64'
+
+    def set_input(self):
+        self.input_shape = [0, 4]
+        self.index_shape = [0, 3]
+        self.input_np = np.random.random(self.input_shape).astype(
+            self.input_dtype
+        )
+        self.index_np = np.random.random(self.index_shape).astype(
+            self.index_dtype
+        )
+
+    def setUp(self):
+        self.set_mode()
+        self.set_dtype()
+        self.set_input()
+        self.place = (
+            base.CUDAPlace(0)
+            if core.is_compiled_with_cuda()
+            else base.CPUPlace()
+        )
+
+    def test_dygraph(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.input_np)
+        x.stop_gradient = False
+        index = paddle.to_tensor(self.index_np)
+        dy_result = paddle.take(x, index, mode=self.mode)
+        np.testing.assert_allclose(
+            np.take(self.input_np, self.index_np, mode=self.mode),
+            dy_result.numpy(),
+        )
+        loss = paddle.sum(dy_result)
+        loss.backward()
+        np.testing.assert_allclose(x.grad.shape, x.shape)
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -7162,7 +7162,8 @@ def take(`
`7162`	`7162`	`)`
`7163`	`7163`	`elif mode == 'clip':`
`7164`	`7164`	`# 'clip' mode disables indexing with negative numbers.`
`7165`		`- index_1d = clip(index_1d, 0, max_index - 1)`
	`7165`	`+ if max_index > 0: # If max_index is 0, input_1d is 0-size.`
	`7166`	`+ index_1d = clip(index_1d, 0, max_index - 1)`
`7166`	`7167`
`7167`	`7168`	`out = input_1d.index_select(index_1d).reshape(index.shape)`
`7168`	`7169`