[0-size Tensor No.130] Add 0-size Tensor support for logsumexp (#73507)

co63oc · web-flow · commit 993b1edd8b61 · 2025-06-24T09:52:33.000+08:00
* Fix

* Fix

* Fix
diff --git a/paddle/phi/kernels/gpu/logsumexp_kernel.cu b/paddle/phi/kernels/gpu/logsumexp_kernel.cu
@@ -21,6 +21,7 @@
 #include "paddle/phi/kernels/activation_kernel.h"
 #include "paddle/phi/kernels/elementwise_add_kernel.h"
 #include "paddle/phi/kernels/elementwise_subtract_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/activation_functor.h"
 #include "paddle/phi/kernels/funcs/elementwise_base.h"
 #include "paddle/phi/kernels/funcs/transpose_function.cu.h"
@@ -89,6 +90,11 @@ void LogsumexpKernel(const Context& dev_ctx,
                      bool keepdim,
                      bool reduce_all,
                      DenseTensor* out) {
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), -INFINITY, out);
+    return;
+  }
   std::vector<int64_t> axis;
   axis.reserve(axis_in.size());
   std::for_each(axis_in.begin(), axis_in.end(), [&axis](const int& t) {
diff --git a/paddle/phi/kernels/impl/logsumexp_grad_kernel_impl.h b/paddle/phi/kernels/impl/logsumexp_grad_kernel_impl.h
@@ -58,6 +58,10 @@ void LogsumexpGradKernel(const Context& dev_ctx,
                          bool keepdim UNUSED,
                          bool reduce_all,
                          DenseTensor* in_grad) {
+  if (in_grad && in_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(in_grad);
+    return;
+  }
   std::vector<int64_t> axis;
   axis.reserve(axis_in.size());
   std::for_each(axis_in.begin(), axis_in.end(), [&axis](const int& t) {
diff --git a/paddle/phi/kernels/impl/logsumexp_kernel_impl.h b/paddle/phi/kernels/impl/logsumexp_kernel_impl.h
@@ -17,11 +17,11 @@
 #include <vector>
 
 #include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
 #include "paddle/phi/kernels/logsumexp_kernel.h"
-
 namespace phi {
 
 #define HANDLE_DIM(NDIM, RDIM)                                         \
@@ -67,6 +67,11 @@ void LogsumexpKernel(const Context& dev_ctx,
                      bool keepdim,
                      bool reduce_all,
                      DenseTensor* out) {
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), -INFINITY, out);
+    return;
+  }
   std::vector<int64_t> axis;
   axis.reserve(axis_in.size());
   std::for_each(axis_in.begin(), axis_in.end(), [&axis](const int& t) {
diff --git a/paddle/phi/kernels/xpu/logsumexp_grad_kernel.cc b/paddle/phi/kernels/xpu/logsumexp_grad_kernel.cc
@@ -27,6 +27,10 @@ void LogsumexpGradKernel(const Context& dev_ctx,
                          bool keepdim,
                          bool reduce_all,
                          DenseTensor* dx) {
+  if (dx && dx->numel() == 0) {
+    dev_ctx.template Alloc<T>(dx);
+    return;
+  }
   using XPUType = typename XPUTypeTrait<T>::Type;
   xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
   reduce_all = recompute_reduce_all(x, axis_in, reduce_all);
diff --git a/paddle/phi/kernels/xpu/logsumexp_kernel.cc b/paddle/phi/kernels/xpu/logsumexp_kernel.cc
@@ -18,6 +18,7 @@
 #include "paddle/phi/kernels/activation_kernel.h"
 #include "paddle/phi/kernels/elementwise_add_kernel.h"
 #include "paddle/phi/kernels/elementwise_subtract_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/reduce_max_kernel.h"
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
 
@@ -30,6 +31,11 @@ void LogsumexpKernel(const Context& dev_ctx,
                      bool keepdim,
                      bool reduce_all,
                      DenseTensor* out) {
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(out->dims())), -INFINITY, out);
+    return;
+  }
   auto xdim = x.dims();
   for (int i = 0; i < xdim.size(); i++)
     PADDLE_ENFORCE_LT(0,
diff --git a/test/legacy_test/test_logsumexp.py b/test/legacy_test/test_logsumexp.py
@@ -298,19 +298,50 @@ def test_alias(self):
         paddle.enable_static()
 
 
-# Test logsumexp bug
-class TestLogZeroError(unittest.TestCase):
-    def test_errors(self):
-        with paddle.base.dygraph.guard():
+class TestLogsumexp_ZeroSize(OpTest):
+    def setUp(self):
+        self.op_type = 'logsumexp'
+        self.python_api = logsumexp_wrapper
+        self.public_python_api = logsumexp_wrapper
+        self.dtype = 'float64'
+        self.shape = [2, 3, 0]
+        self.axis = [-1]  # out return shape [2, 3], value -inf
+        self.keepdim = False
+        self.reduce_all = False
+        self.set_attrs()
 
-            def test_0_size():
-                array = np.array([], dtype=np.float32)
-                x = paddle.to_tensor(
-                    np.reshape(array, [0, 0, 0]), dtype='float32'
-                )
-                paddle.logsumexp(x, axis=1)
+        np.random.seed(10)
+        x = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
+        out = ref_logsumexp(x, self.axis, self.keepdim, self.reduce_all)
 
-            self.assertRaises(ValueError, test_0_size)
+        self.inputs = {'X': x}
+        self.outputs = {'Out': out}
+        self.attrs = {
+            'axis': self.axis,
+            'keepdim': self.keepdim,
+            'reduce_all': self.reduce_all,
+        }
+
+    def set_attrs(self):
+        pass
+
+    def test_check_output(self):
+        self.check_output(
+            check_pir=True,
+        )
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['X'],
+            ['Out'],
+            check_pir=True,
+        )
+
+
+class TestLogsumexp_ZeroSize2(TestLogsumexp_ZeroSize):
+    def set_attrs(self):
+        self.shape = [2, 3, 0]
+        self.axis = [1]  # out return shape [2, 0]
 
 
 if __name__ == '__main__':
