[0-size Tensor No.344、352] Add 0-size Tensor support for trunc (#72861)

co63oc · web-flow · commit 9821267c8ec1 · 2025-05-27T14:17:55.000+08:00
* Fix

* Fix
diff --git a/paddle/phi/kernels/gpu/trunc_grad_kernel.cu b/paddle/phi/kernels/gpu/trunc_grad_kernel.cu
@@ -34,6 +34,9 @@ void TruncGradKernel(const Context& dev_ctx,
                      DenseTensor* in_grad) {
   const auto* out_grad_data = out_grad.data<T>();
   T* in_grad_data = dev_ctx.template Alloc<T>(in_grad);
+  if (out_grad.numel() == 0) {
+    return;
+  }
 
   int64_t numel = out_grad.numel();
 
diff --git a/paddle/phi/kernels/gpu/trunc_kernel.cu b/paddle/phi/kernels/gpu/trunc_kernel.cu
@@ -71,6 +71,9 @@ void TruncKernel(const Context& dev_ctx,
                  DenseTensor* out) {
   const auto* x_data = x.data<T>();
   auto* out_data = dev_ctx.template Alloc<T>(out);
+  if (x.numel() == 0) {
+    return;
+  }
 
   int64_t numel = x.numel();
   auto config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel);
diff --git a/test/legacy_test/test_trunc_op.py b/test/legacy_test/test_trunc_op.py
@@ -30,13 +30,17 @@ def setUp(self):
         self.python_api = paddle.trunc
         self.public_python_api = paddle.trunc
         self.init_dtype_type()
+        self.init_shape()
         np.random.seed(2021)
-        self.inputs = {'X': np.random.random((20, 20)).astype(self.dtype)}
+        self.inputs = {'X': np.random.random(self.shape).astype(self.dtype)}
         self.outputs = {'Out': (np.trunc(self.inputs['X']))}
 
     def init_dtype_type(self):
         self.dtype = np.float64
 
+    def init_shape(self):
+        self.shape = (20, 20)
+
     def test_check_output(self):
         self.check_output(
             check_pir=True, check_prim_pir=True, check_symbol_infer=False
@@ -70,6 +74,11 @@ def test_check_grad(self):
         pass
 
 
+class TestTruncOp_ZeroSize(TestTruncOp):
+    def init_shape(self):
+        self.shape = (20, 0)
+
+
 class TestTruncAPI(unittest.TestCase):
     def setUp(self):
         self.shape = [20, 20]
