[0-size Tensor No.73] Add 0-size Tensor support for paddle.gammaln API (#72895)

DanielSun11 · web-flow · commit df3428270e55 · 2025-05-28T11:05:56.000+08:00
* fix 0 size Tensor for expand kernel in onednn

* format

* expand_grad support 0 size Tensor

* fix bug and add  unittest

* modify the unittest for onednn kernel

* modify expand unittest to force enable onednn op

* fix gammaln to support 0 size Tensor
diff --git a/paddle/phi/kernels/impl/gammaln_grad_kernel_impl.h b/paddle/phi/kernels/impl/gammaln_grad_kernel_impl.h
@@ -81,6 +81,10 @@ void GammalnGradKernel(const Context& dev_ctx,
                        const DenseTensor& d_out,
                        DenseTensor* d_x) {
   auto numel = d_out.numel();
+  if (d_x && d_x->numel() == 0) {
+    dev_ctx.template Alloc<T>(d_x);
+    return;
+  }
   auto* dout_data = d_out.data<T>();
   auto* x_data = x.data<T>();
   auto* dx_data =
diff --git a/paddle/phi/kernels/impl/gammaln_kernel_impl.h b/paddle/phi/kernels/impl/gammaln_kernel_impl.h
@@ -42,6 +42,9 @@ void GammalnKernel(const Context& dev_ctx,
   auto numel = x.numel();
   auto* x_data = x.data<T>();
   auto* out_data = dev_ctx.template Alloc<T>(out);
+  if (numel == 0) {
+    return;
+  }
   phi::funcs::ForRange<Context> for_range(dev_ctx, numel);
   GammalnFunctor<T> functor(x_data, out_data, numel);
   for_range(functor);
diff --git a/test/legacy_test/test_gammaln_op.py b/test/legacy_test/test_gammaln_op.py
@@ -35,12 +35,15 @@ def setUp(self):
         self.op_type = 'gammaln'
         self.python_api = paddle.gammaln
         self.init_dtype_type()
-        self.shape = (3, 40)
+        self.init_shape()
         self.x = np.random.random(self.shape).astype(self.dtype) + 1
         self.inputs = {'x': self.x}
         out = ref_gammaln(self.x)
         self.outputs = {'out': out}
 
+    def init_shape(self):
+        self.shape = (3, 40)
+
     def init_dtype_type(self):
         self.dtype = np.float64
 
@@ -51,6 +54,21 @@ def test_check_grad(self):
         self.check_grad(['x'], 'out', check_pir=True)
 
 
+class TestGammalnOpZeroSize(TestGammalnOp):
+    def init_shape(self):
+        self.shape = (0, 3, 40)
+
+
+class TestGammalnOpZeroSize1(TestGammalnOp):
+    def init_shape(self):
+        self.shape = (10, 3, 0, 1)
+
+
+class TestGammalnOpZeroSize2(TestGammalnOp):
+    def init_shape(self):
+        self.shape = (10, 0)
+
+
 class TestGammalnOpFp32(TestGammalnOp):
     def init_dtype_type(self):
         self.dtype = np.float32
