[0-size Tensor No.343, 350, 351] Add 0-size Tensor support for paddle.tril paddle.Tensor.tril paddle.triu API. (#72814)

crashbussy · web-flow · commit 263f2e013f83 · 2025-05-28T10:58:06.000+08:00
* Create test_tril_op.py

增加对tril的单测

* Create test_triu_op.py

增加对triu的单测

* Update test_tril_triu_op.py

增加了对0size 的支持

* Create test_xpu_tril_triu_0size_op.py

xpu的单测

* Rename test_xpu_tril_triu_0size_op.py to test_tril_triu_xpu_0size_op.py

* Delete test/legacy_test/test_triu_op.py

* Delete test/legacy_test/test_tril_op.py

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Fix code style issues

* Update tril_triu_kernel_impl.h

* Update tril_triu_grad_kernel_impl.h

* Update tril_triu_kernel_impl.h

* Update tril_triu_grad_kernel_impl.h

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Update test_tril_triu_op_xpu.py

* Delete test/xpu/test_tril_triu_xpu_0size_op.py

* Update tril_triu_kernel_impl.h

* Update test_tril_triu_op_xpu.py

* Update test_tril_triu_op.py

* Update tril_triu_grad_kernel_impl.h

* Update test_tril_triu_op.py

* Update test_tril_triu_op_xpu.py

* Fix code style issues with pre-commit

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* precommitfix

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Update tril_triu_grad_kernel_impl.h

* Update test_tril_triu_op.py

* Update tril_triu_grad_kernel_impl.h

* Update test_tril_triu_op.py

* Update tril_triu_grad_kernel_impl.h

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* Update test_tril_triu_op.py

* precommit fix

* Update test_tril_triu_op.py

* precommit fix

* Update tril_triu_kernel.cc

* Update tril_triu_grad_kernel.cc
diff --git a/paddle/phi/kernels/impl/tril_triu_grad_kernel_impl.h b/paddle/phi/kernels/impl/tril_triu_grad_kernel_impl.h
@@ -45,6 +45,11 @@ void TrilGradKernel(const Context& ctx,
                     const DenseTensor& out_grad,
                     int diagonal,
                     DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
+
   TrilTriuGradKernel<T, Context>(ctx, out_grad, diagonal, true, x_grad);
 }
 
@@ -53,6 +58,11 @@ void TriuGradKernel(const Context& ctx,
                     const DenseTensor& out_grad,
                     int diagonal,
                     DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
+
   TrilTriuGradKernel<T, Context>(ctx, out_grad, diagonal, false, x_grad);
 }
 
diff --git a/paddle/phi/kernels/impl/tril_triu_kernel_impl.h b/paddle/phi/kernels/impl/tril_triu_kernel_impl.h
@@ -44,6 +44,11 @@ void TrilKernel(const Context& ctx,
                 const DenseTensor& x,
                 int diagonal,
                 DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
+
   TrilTriuKernel<T, Context>(ctx, x, diagonal, true, out);
 }
 
@@ -52,6 +57,11 @@ void TriuKernel(const Context& ctx,
                 const DenseTensor& x,
                 int diagonal,
                 DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
+
   TrilTriuKernel<T, Context>(ctx, x, diagonal, false, out);
 }
 
diff --git a/paddle/phi/kernels/xpu/tril_triu_grad_kernel.cc b/paddle/phi/kernels/xpu/tril_triu_grad_kernel.cc
@@ -51,6 +51,11 @@ void TrilGradKernel(const Context& ctx,
                     const DenseTensor& out_grad,
                     int diagonal,
                     DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
+
   TrilTriuGradKernel<T, Context>(ctx, out_grad, diagonal, true, x_grad);
 }
 
@@ -59,6 +64,11 @@ void TriuGradKernel(const Context& ctx,
                     const DenseTensor& out_grad,
                     int diagonal,
                     DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    return;
+  }
+
   TrilTriuGradKernel<T, Context>(ctx, out_grad, diagonal, false, x_grad);
 }
 
diff --git a/paddle/phi/kernels/xpu/tril_triu_kernel.cc b/paddle/phi/kernels/xpu/tril_triu_kernel.cc
@@ -51,6 +51,11 @@ void TrilKernel(const Context& ctx,
                 const DenseTensor& x,
                 int diagonal,
                 DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
+
   TrilTriuKernel<T, Context>(ctx, x, diagonal, true, out);
 }
 
@@ -59,6 +64,11 @@ void TriuKernel(const Context& ctx,
                 const DenseTensor& x,
                 int diagonal,
                 DenseTensor* out) {
+  if (out && out->numel() == 0) {
+    ctx.template Alloc<T>(out);
+    return;
+  }
+
   TrilTriuKernel<T, Context>(ctx, x, diagonal, false, out);
 }
 
diff --git a/test/legacy_test/test_tril_triu_op.py b/test/legacy_test/test_tril_triu_op.py
@@ -327,5 +327,60 @@ def test_base_api(self):
                 )
 
 
+class TestTrilZeroSizeShape(TrilTriuOpDefaultTest):
+    def initTestCase(self):
+        self.real_op_type = 'tril'
+        self.diagonal = 0
+        self.X = np.random.rand(0, 3, 9, 4).astype(np.float64)
+
+
+class TestTriuZeroSizeShape(TrilTriuOpDefaultTest):
+    def initTestCase(self):
+        self.real_op_type = 'triu'
+        self.diagonal = 0
+        self.X = np.random.rand(0, 3, 9, 4).astype(np.float64)
+
+
+class TestTrilTriu_ZeroDimGrad(OpTest):
+    def setUp(self):
+        self.op_type = 'tril_triu'
+        self.real_op_type = 'tril'
+        self.diagonal = 0
+        self.inputs = {'X': np.random.randn(0, 3, 9, 4).astype('float64')}
+        self.attrs = {
+            'diagonal': self.diagonal,
+            'lower': True if self.real_op_type == 'tril' else False,
+        }
+        self.outputs = {
+            'Out': getattr(np, self.real_op_type)(
+                self.inputs['X'], self.diagonal
+            )
+        }
+        self.python_api = (
+            paddle.tril if self.real_op_type == 'tril' else paddle.triu
+        )
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True)
+
+
+class TestTrilTriu_ZeroDimGrad_Triu(OpTest):
+    def setUp(self):
+        self.op_type = 'tril_triu'
+        self.real_op_type = 'triu'
+        self.diagonal = 0
+        self.inputs = {'X': np.random.randn(0, 3, 9, 4).astype('float64')}
+        self.attrs = {'diagonal': self.diagonal, 'lower': False}
+        self.outputs = {
+            'Out': getattr(np, self.real_op_type)(
+                self.inputs['X'], self.diagonal
+            )
+        }
+        self.python_api = paddle.triu
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', check_pir=True)
+
+
 if __name__ == '__main__':
     unittest.main()
