[0-size Tensor No.33、35、49、50、74、75、253、263、274、335、361] Add 0-size Tensor support for some API which using split or concat. (#72988)

DanielSun11 · web-flow · commit 98db28aae05d · 2025-05-30T16:26:59.000+08:00
* fix split

* fix concat and add unittest

* add unittest case

* fix unittest case
diff --git a/paddle/phi/kernels/cpu/concat_kernel.cc b/paddle/phi/kernels/cpu/concat_kernel.cc
@@ -45,7 +45,9 @@ void ConcatKernel(const Context& dev_ctx,
   phi::DDim out_dims = phi::funcs::ComputeAndCheckShape(true, x_dims, axis);
   out->Resize(out_dims);
   dev_ctx.template Alloc<T>(out);
-
+  if (out->numel() == 0) {
+    return;
+  }
   // If axis is 0, the lod of the output is not the same as inputs.
   if (axis == 0 && !x[0]->lod().empty()) {
     size_t lod_size_0 = x[0]->lod().size();
diff --git a/paddle/phi/kernels/impl/concat_grad_kernel_impl.h b/paddle/phi/kernels/impl/concat_grad_kernel_impl.h
@@ -14,6 +14,7 @@
 #pragma once
 
 #include "paddle/phi/kernels/concat_grad_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
 #include "paddle/phi/kernels/funcs/concat_funcs.h"
 #include "paddle/phi/kernels/funcs/strided_memcpy.h"
@@ -52,6 +53,11 @@ void ConcatGradKernel(const Context& dev_ctx,
       outputs.push_back(nullptr);
     }
   }
+  // if the out_grad.numel() == 0 ,the all x and x_grad must be zero size
+  // tensor, so just return
+  if (out_grad.numel() == 0) {
+    return;
+  }
 
   // Sometimes direct copies will be faster, this maybe need deeply analysis.
   if (axis == 0 && outs.size() < 10) {
diff --git a/paddle/phi/kernels/xpu/concat_grad_kernel.cc b/paddle/phi/kernels/xpu/concat_grad_kernel.cc
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/axis_utils.h"
 #include "paddle/phi/kernels/funcs/concat_funcs.h"
 
@@ -52,6 +53,11 @@ void ConcatGradKernel(const Context& dev_ctx,
       ptrs[j] = reinterpret_cast<XPUType*>(outs[j]->data<T>());
     }
   }
+  // if the out_grad.numel() == 0 ,the all x and x_grad must be zero size
+  // tensor, so just return
+  if (out_grad.numel() == 0) {
+    return;
+  }
   PADDLE_ENFORCE_GE(axis,
                     0,
                     common::errors::InvalidArgument(
diff --git a/paddle/phi/kernels/xpu/concat_kernel.cc b/paddle/phi/kernels/xpu/concat_kernel.cc
@@ -57,6 +57,9 @@ void ConcatKernel(const Context& dev_ctx,
   phi::DDim out_dims = phi::funcs::ComputeAndCheckShape(true, x_dims, axis);
   out->Resize(out_dims);
   dev_ctx.template Alloc<T>(out);
+  if (out->numel() == 0) {
+    return;
+  }
 
   // If axis is 0, the lod of the output is not the same as inputs.
 
diff --git a/test/legacy_test/test_concat_op.py b/test/legacy_test/test_concat_op.py
@@ -1066,6 +1066,46 @@ def test_empty_inputs_static(self):
                 paddle.concat([], axis=0)
 
 
+class TestConcatOpZeroSize1(TestConcatOp):
+    def init_test_data(self):
+        self.x0 = np.random.random((2, 0, 4, 5)).astype(self.dtype)
+        self.x1 = np.random.random((2, 3, 4, 5)).astype(self.dtype)
+        self.x2 = np.random.random((2, 3, 4, 5)).astype(self.dtype)
+        self.axis = 1
+
+
+class TestConcatOpZeroSize2(TestConcatOp):
+    def init_test_data(self):
+        self.x0 = np.random.random((2, 0, 1, 5)).astype(self.dtype)
+        self.x1 = np.random.random((2, 0, 2, 5)).astype(self.dtype)
+        self.x2 = np.random.random((2, 0, 4, 5)).astype(self.dtype)
+        self.axis = 2
+
+
+class TestConcatOpZeroSize3(TestConcatOp):
+    def init_test_data(self):
+        self.x0 = np.random.random((0, 0, 0, 0)).astype(self.dtype)
+        self.x1 = np.random.random((0, 0, 0, 0)).astype(self.dtype)
+        self.x2 = np.random.random((0, 0, 0, 0)).astype(self.dtype)
+        self.axis = 2
+
+
+class TestConcatOpZeroSize4(TestConcatOp):
+    def init_test_data(self):
+        self.x0 = np.random.random((0, 1, 2, 3)).astype(self.dtype)
+        self.x1 = np.random.random((0, 1, 2, 3)).astype(self.dtype)
+        self.x2 = np.random.random((0, 1, 2, 3)).astype(self.dtype)
+        self.axis = 2
+
+
+class TestConcatOpZeroSize5(TestConcatOp):
+    def init_test_data(self):
+        self.x0 = np.random.random((0, 1, 2, 3)).astype(self.dtype)
+        self.x1 = np.random.random((0, 1, 2, 3)).astype(self.dtype)
+        self.x2 = np.random.random((0, 1, 2, 3)).astype(self.dtype)
+        self.axis = 2
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
diff --git a/test/legacy_test/test_split_op.py b/test/legacy_test/test_split_op.py
@@ -142,11 +142,13 @@ def setUp(self):
             'X': self.x,
             'AxisTensor': np.array([self.axis]).astype("int32"),
         }
-        self.attrs = {'sections': self.sections, 'num': self.num}
-
+        self.init_attrs()
         out = np.split(self.x, self.indices_or_sections, self.axis)
         self.outputs = {'Out': [(f'out{i}', out[i]) for i in range(len(out))]}
 
+    def init_attrs(self):
+        self.attrs = {'sections': self.sections, 'num': self.num}
+
     def init_data(self):
         self.x = np.random.random((4, 5, 6)).astype(self.dtype)
         self.axis = 2
@@ -167,6 +169,47 @@ def test_check_grad(self):
         self.check_grad(['X'], ['out0', 'out1', 'out2'], check_pir=True)
 
 
+class TestSplitOpZeroSize(TestSplitOp_AxisTensor):
+    def init_data(self):
+        self.x = np.random.random((0, 1, 6)).astype(self.dtype)
+        self.axis = 2
+        self.sections = [2, 2, 2]
+        self.indices_or_sections = 3
+
+    def init_attrs(self):
+        self.attrs = {'sections': self.sections, 'axis': self.axis}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
+    def test_check_grad(self):
+        self.check_grad(['X'], ['out0', 'out1', 'out2'], check_pir=True)
+
+
+class TestSplitOpZeroSize1(TestSplitOpZeroSize):
+    def init_data(self):
+        self.x = np.random.random((8, 0, 9)).astype(self.dtype)
+        self.axis = 2
+        self.sections = [1, 4, 4]
+        self.indices_or_sections = [1, 5]
+
+
+class TestSplitOpZeroSize2(TestSplitOpZeroSize):
+    def init_data(self):
+        self.x = np.random.random((5, 0, 12)).astype(self.dtype)
+        self.axis = 2
+        self.sections = [6, 0, 6]
+        self.indices_or_sections = [6, 6]
+
+
+class TestSplitOpZeroSize3(TestSplitOpZeroSize):
+    def init_data(self):
+        self.x = np.random.random((5, 0, 12)).astype(self.dtype)
+        self.axis = 1
+        self.sections = [6, 0, 6]
+        self.indices_or_sections = [6, 6]
+
+
 # attr(sections) is list containing Tensor
 class TestSplitOp_SectionsTensor(OpTest):
     def setUp(self):
@@ -745,6 +788,67 @@ def test_axis_input_empty_section(self):
         np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
 
 
+class API_TestSplitZeroSize(unittest.TestCase):
+    def test_case1(self):
+        with base.dygraph.guard():
+            input_1 = np.random.random([3, 0, 6]).astype("float32")
+            input = paddle.to_tensor(input_1)
+            x0, x1, x2 = paddle.split(input, num_or_sections=[3, 0, 3], axis=-1)
+            x0_out = x0.numpy()
+            x1_out = x1.numpy()
+            x2_out = x2.numpy()
+            ex_x0, ex_x1, ex_x2 = np.split(
+                input_1,
+                [3, 3],
+                axis=-1,
+            )
+        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
+        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
+        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
+
+    def test_case2(self):
+        with base.dygraph.guard():
+            input_1 = np.random.random([9, 0, 0]).astype("float32")
+            input = paddle.to_tensor(input_1, stop_gradient=False)
+            (
+                x0,
+                x1,
+                x2,
+            ) = paddle.split(input, num_or_sections=3, axis=0)
+            x0_out = x0.numpy()
+            x1_out = x1.numpy()
+            x2_out = x2.numpy()
+            ex_x0, ex_x1, ex_x2 = np.split(
+                input_1,
+                3,
+                axis=0,
+            )
+        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
+        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
+        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
+
+    def test_case3(self):
+        with base.dygraph.guard():
+            input_1 = np.random.random([9, 0, 0]).astype("float32")
+            input = paddle.to_tensor(input_1)
+            (
+                x0,
+                x1,
+                x2,
+            ) = paddle.split(input, num_or_sections=3, axis=1)
+            x0_out = x0.numpy()
+            x1_out = x1.numpy()
+            x2_out = x2.numpy()
+            ex_x0, ex_x1, ex_x2 = np.split(
+                input_1,
+                3,
+                axis=1,
+            )
+        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
+        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
+        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
