【Hackathon 6th No.9】Update stack op to support zero sized tensor -part (#64977)

NKNaN · web-flow · commit f56c21ff799a · 2024-06-27T17:35:54.000+08:00
* update stack op

* update dirichlet

* update stack test

* add docs example

* fix dirichlet

* rerun ci
diff --git a/paddle/phi/kernels/cpu/stack_grad_kernel.cc b/paddle/phi/kernels/cpu/stack_grad_kernel.cc
@@ -37,6 +37,16 @@ void StackGradKernel(const Context& dev_ctx,
     }
   }
   auto dy_data = out.data<T>();
+
+  // zero sized tensor case
+  if (out.numel() == 0) {
+    for (int i = 0; i < n; i++) {
+      auto x_grad_dim = x_grad[i]->dims();
+      x_grad[i]->Resize(x_grad_dim);
+    }
+    return;
+  }
+
   int pre = 1;
   for (int i = 0; i < axis; ++i) pre *= static_cast<int>(out.dims()[i]);
   int total_num = static_cast<int>(out.numel());
diff --git a/paddle/phi/kernels/cpu/stack_kernel.cc b/paddle/phi/kernels/cpu/stack_kernel.cc
@@ -28,10 +28,18 @@ void StackKernel(const Context& dev_ctx,
 
   auto x_dims = x[0]->dims();
   for (int i = 0; i < x_dims.size(); i++) {
-    PADDLE_ENFORCE_GT(x_dims[i],
-                      0,
-                      phi::errors::InvalidArgument(
-                          "The dims of Input(X) should be greater than 0"));
+    PADDLE_ENFORCE_GE(
+        x_dims[i],
+        0,
+        phi::errors::InvalidArgument(
+            "The dims of Input(X) should be greater than or equal to 0"));
+  }
+  // zero sized tensor case
+  if (x[0]->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    auto out_dims = out->dims();
+    out->Resize(out_dims);
+    return;
   }
 
   int n = static_cast<int>(x.size());
diff --git a/paddle/phi/kernels/funcs/stack_and_unstack.h b/paddle/phi/kernels/funcs/stack_and_unstack.h
@@ -77,6 +77,13 @@ void StackRawKernel(const Context& ctx,
   if (axis < 0) axis += (x[0]->dims().size() + 1);
   int num = static_cast<int>(x.size());
 
+  // zero sized tensor case
+  if (x[0]->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    auto out_dims = out->dims();
+    out->Resize(out_dims);
+    return;
+  }
   // Split x dim from axis to matrix of shape [x_row, x_col], and the output
   // tensor's shape is [x_row, out_col].
   int64_t x_row = 1, x_row_bak = 1;
@@ -251,6 +258,15 @@ void UnStackRawKernel(const Context& ctx,
   // Input tensor is splited to split_dim tensors along split_dim dimension.
   int64_t split_dim = x_dims[axis];
 
+  // zero sized tensor case
+  if (x.numel() == 0) {
+    for (int i = 0; i < split_dim; i++) {
+      ctx.template Alloc<T>((*outs)[i]);
+      auto x_grad_dim = (*outs)[i]->dims();
+      (*outs)[i]->Resize(x_grad_dim);
+    }
+    return;
+  }
   // Treat outs[i] as [out_row, out_col], and x as [out_row, split_dim,
   // out_col].
   int64_t out_row = 1;
diff --git a/python/paddle/distribution/dirichlet.py b/python/paddle/distribution/dirichlet.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+
 import paddle
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.layer_helper import LayerHelper
@@ -69,7 +71,8 @@ class Dirichlet(exponential_family.ExponentialFamily):
     """
 
     def __init__(self, concentration):
-        if concentration.dim() < 1:
+        if concentration.dim() < 1 or math.prod(concentration.shape) == 0:
+            # 0-dim tensor or 0-sized tensor is invalid
             raise ValueError(
                 "`concentration` parameter must be at least one dimensional"
             )
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -2048,6 +2048,8 @@ def stack(
     tensor is [N, A, B]; if ``axis == 1``, the shape of stacked
     tensor is [A, N, B], etc.
 
+    It also supports the operation with zero-size tensors which contain 0 in their shape.
+    See the examples below.
 
     .. code-block:: text
 
@@ -2091,6 +2093,39 @@ def stack(
                           [3.0, 4.0]
                           [5.0, 6.0] ] ]
 
+
+        Case 3:
+
+            Input:
+                x[0].shape = [0, 1, 2]
+                x[0].data = []
+                x[1].shape = [0, 1, 2]
+                x[1].data = []
+
+            Attrs:
+                axis = 0
+
+            Output:
+                Out.shape = [2, 0, 1, 2]
+                Out.data = []
+
+
+        Case 4:
+
+            Input:
+                x[0].shape = [0, 1, 2]
+                x[0].data = []
+                x[1].shape = [0, 1, 2]
+                x[1].data = []
+
+            Attrs:
+                axis = 1
+
+            Output:
+                Out.shape = [0, 2, 1, 2]
+                Out.data = []
+
+
     Args:
         x (list[Tensor]|tuple[Tensor]): Input ``x`` can be a ``list`` or ``tuple`` of tensors, the Tensors in ``x``
                                      must be of the same shape and dtype. Supported data types: float32, float64, int32, int64.
@@ -2128,6 +2163,25 @@ def stack(
             [[[1., 2.],
               [3., 4.],
               [5., 6.]]])
+
+            >>> # zero-size tensors
+            >>> x1 = paddle.ones([0, 1, 2])
+            >>> x2 = paddle.ones([0, 1, 2])
+
+            >>> out = paddle.stack([x1, x2], axis=0)
+            >>> print(out.shape)
+            [2, 0, 1, 2]
+            >>> print(out)
+            Tensor(shape=[2, 0, 1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[],
+             []])
+
+            >>> out = paddle.stack([x1, x2], axis=1)
+            >>> print(out.shape)
+            [0, 2, 1, 2]
+            >>> print(out)
+            Tensor(shape=[0, 2, 1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [])
     """
     axis = 0 if axis is None else axis
 
diff --git a/test/legacy_test/test_stack_op.py b/test/legacy_test/test_stack_op.py
@@ -452,5 +452,97 @@ def test_stack_triple_grad(self):
         paddle.base.core.set_prim_eager_enabled(False)
 
 
+class TestStackAPI_ZeroSizedTensor(unittest.TestCase):
+    def test_dygraph_cpu(self):
+        place = base.CPUPlace()
+        paddle.disable_static(place)
+
+        x1 = paddle.ones([1, 0])
+        x2 = paddle.ones([1, 0])
+        x1.stop_gradient = False
+        x2.stop_gradient = False
+        out = paddle.stack([x1, x2])
+        out.retain_grads()
+        out.backward()
+
+        np.testing.assert_equal(out.shape, [2, 1, 0])
+        np.testing.assert_equal(x1.grad, None)
+        np.testing.assert_equal(x2.grad, None)
+        np.testing.assert_equal(out, np.ones([2, 1, 0]))
+
+        paddle.enable_static()
+
+    def test_dygraph_gpu(self):
+        if base.is_compiled_with_cuda():
+            place = base.CUDAPlace(0)
+            paddle.disable_static(place)
+
+            x1 = paddle.ones([1, 0])
+            x2 = paddle.ones([1, 0])
+            x1.stop_gradient = False
+            x2.stop_gradient = False
+            out = paddle.stack([x1, x2])
+            out.retain_grads()
+            out.backward()
+
+            np.testing.assert_equal(out.shape, [2, 1, 0])
+            np.testing.assert_equal(x1.grad, None)
+            np.testing.assert_equal(x2.grad, None)
+            np.testing.assert_equal(out, np.ones([2, 1, 0]))
+
+            paddle.enable_static()
+
+    @test_with_pir_api
+    def test_static_cpu(self):
+        paddle.enable_static()
+        place = base.CPUPlace()
+        exe = base.Executor(place)
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            data1 = paddle.static.data('data1', shape=[0, 2], dtype='float64')
+            data2 = paddle.static.data('data2', shape=[0, 2], dtype='float64')
+            data3 = paddle.static.data('data3', shape=[0, 2], dtype='float64')
+            result_stack = paddle.stack([data1, data2, data3], axis=0)
+            input1 = np.ones([0, 2]).astype('float64')
+            input2 = np.ones([0, 2]).astype('float64')
+            input3 = np.ones([0, 2]).astype('float64')
+            (result,) = exe.run(
+                feed={"data1": input1, "data2": input2, "data3": input3},
+                fetch_list=[result_stack],
+            )
+            expected_result = np.stack([input1, input2, input3], axis=0)
+            np.testing.assert_equal(expected_result, result)
+
+    @test_with_pir_api
+    def test_static_gpu(self):
+        if base.is_compiled_with_cuda():
+            paddle.enable_static()
+            place = base.CUDAPlace(0)
+            exe = base.Executor(place)
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                data1 = paddle.static.data(
+                    'data1', shape=[0, 2], dtype='float64'
+                )
+                data2 = paddle.static.data(
+                    'data2', shape=[0, 2], dtype='float64'
+                )
+                data3 = paddle.static.data(
+                    'data3', shape=[0, 2], dtype='float64'
+                )
+                result_stack = paddle.stack([data1, data2, data3], axis=0)
+                input1 = np.ones([0, 2]).astype('float64')
+                input2 = np.ones([0, 2]).astype('float64')
+                input3 = np.ones([0, 2]).astype('float64')
+                (result,) = exe.run(
+                    feed={"data1": input1, "data2": input2, "data3": input3},
+                    fetch_list=[result_stack],
+                )
+                expected_result = np.stack([input1, input2, input3], axis=0)
+                np.testing.assert_equal(expected_result, result)
+
+
 if __name__ == '__main__':
     unittest.main()
