[0-size Tensor No.59、62、250] Add 0-size Tensor support for fftshift etc (#72809)

co63oc · web-flow · commit 9788095ba1fc · 2025-05-28T10:55:58.000+08:00
* Fix

* Fix

* Fix

* Fix
diff --git a/paddle/phi/kernels/cpu/roll_grad_kernel.cc b/paddle/phi/kernels/cpu/roll_grad_kernel.cc
@@ -27,6 +27,10 @@ void RollGradKernel(const Context& dev_ctx,
                     const IntArray& shifts,
                     const std::vector<int64_t>& axis,
                     DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   std::vector<T> out_vec;
   phi::TensorToVector(out_grad, dev_ctx, &out_vec);
 
diff --git a/paddle/phi/kernels/gpu/roll_grad_kernel.cu b/paddle/phi/kernels/gpu/roll_grad_kernel.cu
@@ -30,6 +30,10 @@ void RollGradKernel(const Context& dev_ctx,
                     const std::vector<int64_t>& axis,
                     DenseTensor* x_grad) {
   auto* out_grad_data = out_grad.data<T>();
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
 
   auto shifts_data = shifts.GetData();
diff --git a/paddle/phi/kernels/xpu/roll_grad_kernel.cc b/paddle/phi/kernels/xpu/roll_grad_kernel.cc
@@ -28,6 +28,10 @@ void RollGradKernel(const Context& dev_ctx,
                     DenseTensor* x_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto shifts_data = shifts.GetData();
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   dev_ctx.template Alloc<T>(x_grad);
   DDim input_dim = x.dims();
   std::vector<int64_t> xshape;
diff --git a/paddle/phi/kernels/xpu/roll_kernel.cc b/paddle/phi/kernels/xpu/roll_kernel.cc
@@ -26,6 +26,10 @@ void RollKernel(const Context& dev_ctx,
                 DenseTensor* out) {
   using XPUType = typename XPUTypeTrait<T>::Type;
   auto shifts_data = shifts.GetData();
+  if (out && out->numel() == 0) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   dev_ctx.template Alloc<T>(out);
   DDim input_dim = x.dims();
   std::vector<int64_t> xshape;
diff --git a/test/fft/test_fft.py b/test/fft/test_fft.py
@@ -16,6 +16,7 @@
 import unittest
 
 import numpy as np
+import scipy
 import scipy.fft
 
 import paddle
@@ -1934,5 +1935,79 @@ def test_ifftshift(self):
             )
 
 
+@place(DEVICES)
+@parameterize(
+    (TEST_CASE_NAME, 'x', 'axes', 'dtype'),
+    [
+        ('test_1d', np.random.randn(0), (0,), 'float64'),
+        (
+            'test_2d_odd_with_all_axes',
+            np.random.randn(5, 0) + 1j * np.random.randn(5, 0),
+            None,
+            'complex128',
+        ),
+    ],
+)
+class TestFftShift_ZeroSize(unittest.TestCase):
+    def test_fftshift(self):
+        """Test fftshift with norm condition"""
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.fftshift(self.x, self.axes),
+                paddle.fft.fftshift(
+                    paddle.to_tensor(self.x), self.axes
+                ).numpy(),
+                rtol=RTOL.get(str(self.x.dtype)),
+                atol=ATOL.get(str(self.x.dtype)),
+            )
+
+    def test_grad_shape(self):
+        with paddle.base.dygraph.guard(self.place):
+            x = paddle.to_tensor(self.x, stop_gradient=False)
+            y = paddle.fft.fftshift(x, self.axes)
+            loss = paddle.sum(y)
+            loss.backward()
+            np.testing.assert_equal(
+                x.grad.shape, self.x.shape, "Grad shape mismatch"
+            )
+
+
+@place(DEVICES)
+@parameterize(
+    (TEST_CASE_NAME, 'x', 'axes'),
+    [
+        ('test_1d', np.random.randn(0), (0,), 'float64'),
+        (
+            'test_2d_odd_with_all_axes',
+            np.random.randn(5, 0) + 1j * np.random.randn(5, 0),
+            None,
+            'complex128',
+        ),
+    ],
+)
+class TestIfftShift_ZeroSize(unittest.TestCase):
+    def test_ifftshift(self):
+        """Test ifftshift with norm condition"""
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.ifftshift(self.x, self.axes),
+                paddle.fft.ifftshift(
+                    paddle.to_tensor(self.x), self.axes
+                ).numpy(),
+                rtol=RTOL.get(str(self.x.dtype)),
+                atol=ATOL.get(str(self.x.dtype)),
+            )
+
+    def test_grad_shape(self):
+        with paddle.base.dygraph.guard(self.place):
+            x = paddle.to_tensor(self.x, stop_gradient=False)
+            y = paddle.fft.ifftshift(x, self.axes)
+            loss = paddle.sum(y)
+            loss.backward()
+            np.testing.assert_equal(
+                x.grad.shape, self.x.shape, "Grad shape mismatch"
+            )
+
+
 if __name__ == '__main__':
     unittest.main()
