[0-size Tensor No.57-58、60-61、63-66] Add 0-size Tensor support for fft2 (#73042)

co63oc · web-flow · commit 89a443e0fbb8 · 2025-06-04T19:50:04.000+08:00
* Fix

* Fix

* Fix

* Fix

* Fix
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -1470,7 +1470,7 @@ void FFTC2CInferMeta(const MetaTensor& x,
   if (config.is_runtime) {
     const phi::DDim x_dim = x.dims();
     for (auto axis : axes) {
-      PADDLE_ENFORCE_GT(x_dim[axis],
+      PADDLE_ENFORCE_GE(x_dim[axis],
                         0,
                         common::errors::InvalidArgument(
                             "Invalid fft n-point (%d).", x_dim[axis]));
@@ -1497,7 +1497,7 @@ void FFTC2RInferMeta(const MetaTensor& x,
   if (config.is_runtime) {
     size_t signal_dims = axes.size();
     for (size_t i = 0; i < signal_dims - 1; i++) {
-      PADDLE_ENFORCE_GT(x_dim[axes[i]],
+      PADDLE_ENFORCE_GE(x_dim[axes[i]],
                         0,
                         common::errors::InvalidArgument(
                             "Invalid fft n-point (%d).", x_dim[axes[i]]));
@@ -1513,7 +1513,7 @@ void FFTC2RInferMeta(const MetaTensor& x,
   } else if (config.is_runtime) {
     const int64_t input_last_dim_size = x_dim[last_fft_axis];
     const int64_t fft_n_point = (input_last_dim_size - 1) * 2;
-    PADDLE_ENFORCE_GT(fft_n_point,
+    PADDLE_ENFORCE_GE(fft_n_point,
                       0,
                       common::errors::InvalidArgument(
                           "Invalid fft n-point (%d).", fft_n_point));
@@ -1542,7 +1542,7 @@ void FFTR2CInferMeta(const MetaTensor& x,
   // they might be -1 to indicate unknown size ar compile time
   if (config.is_runtime) {
     for (auto axis : axes) {
-      PADDLE_ENFORCE_GT(x_dim[axis],
+      PADDLE_ENFORCE_GE(x_dim[axis],
                         0,
                         common::errors::InvalidArgument(
                             "Invalid fft n-point (%d).", x_dim[axis]));
diff --git a/paddle/phi/kernels/impl/fft_grad_kernel_impl.h b/paddle/phi/kernels/impl/fft_grad_kernel_impl.h
@@ -38,6 +38,9 @@ void FFTC2CGradKernel(const Context& ctx,
                       bool forward,
                       DenseTensor* x_grad) {
   ctx.template Alloc<T>(x_grad);
+  if (x_grad && x_grad->numel() == 0) {
+    return;
+  }
   auto norm_type = funcs::get_norm_from_string(normalization, forward);
   funcs::FFTC2CFunctor<Context, T, T> fft_c2c_func;
   fft_c2c_func(ctx, out_grad, x_grad, axes, norm_type, !forward);
@@ -55,6 +58,10 @@ void FFTR2CGradKernel(const Context& ctx,
   using R = typename T::value_type;
   DenseTensor complex_x_grad = EmptyLike<T>(ctx, x);
   ctx.template Alloc<R>(x_grad);
+  if (x_grad && x_grad->numel() == 0) {
+    return;
+  }
+
   auto norm_type = funcs::get_norm_from_string(normalization, forward);
   funcs::FFTC2CFunctor<Context, T, T> fft_c2c_func;
 
@@ -85,6 +92,10 @@ void FFTC2RGradKernel(const Context& ctx,
                       DenseTensor* x_grad) {
   using C = phi::dtype::complex<T>;
   ctx.template Alloc<C>(x_grad);
+  if (x_grad && x_grad->numel() == 0) {
+    return;
+  }
+
   auto norm_type = funcs::get_norm_from_string(normalization, forward);
 
   funcs::FFTR2CFunctor<Context, T, C> fft_r2c_func;
diff --git a/paddle/phi/kernels/impl/fft_kernel_impl.h b/paddle/phi/kernels/impl/fft_kernel_impl.h
@@ -21,9 +21,9 @@
 #include "paddle/common/ddim.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/empty_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/fft.h"
 #include "paddle/phi/kernels/funcs/fft_fill_conj.h"
-
 namespace phi {
 template <typename T, typename Context>
 void FFTC2CKernel(const Context& ctx,
@@ -33,6 +33,18 @@ void FFTC2CKernel(const Context& ctx,
                   bool forward,
                   DenseTensor* out) {
   ctx.template Alloc<T>(out);
+  if (x.numel() == 0) {
+    /*
+    This will return 0:
+    >>> scipy.fft.fft2(np.random.random([3, 0, 1, 2]), s=(1, 2), axes=(0, 1),
+norm='backward')
+    array([[[[0.-0.j, 0.-0.j]],
+        [[0.-0.j, 0.-0.j]]]])
+    */
+    phi::Full<T, Context>(
+        ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+    return;
+  }
   const auto norm_type = funcs::get_norm_from_string(normalization, forward);
   funcs::FFTC2CFunctor<Context, T, T> fft_c2c_func;
   fft_c2c_func(ctx, x, out, axes, norm_type, forward);
@@ -48,6 +60,11 @@ void FFTC2RKernel(const Context& ctx,
                   DenseTensor* out) {
   using R = typename T::value_type;  // get real type
   ctx.template Alloc<R>(out);
+  if (x.numel() == 0) {
+    phi::Full<R, Context>(
+        ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+    return;
+  }
   const auto norm_type = funcs::get_norm_from_string(normalization, forward);
   funcs::FFTC2RFunctor<Context, T, R> fft_c2r_func;
   fft_c2r_func(ctx, x, out, axes, norm_type, forward);
@@ -63,6 +80,11 @@ void FFTR2CKernel(const Context& ctx,
                   DenseTensor* out) {
   using C = phi::dtype::complex<T>;
   ctx.template Alloc<C>(out);
+  if (x.numel() == 0) {
+    phi::Full<C, Context>(
+        ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+    return;
+  }
   auto norm_type = funcs::get_norm_from_string(normalization, forward);
   funcs::FFTR2CFunctor<Context, T, C> fft_r2c_func;
 
diff --git a/test/fft/test_fft.py b/test/fft/test_fft.py
@@ -1028,14 +1028,6 @@ def test_hfft2(self):
             'backward',
             ValueError,
         ),
-        (
-            'test_zero_point',
-            np.random.randn(4, 4, 1) + 1j * np.random.randn(4, 4, 1),
-            None,
-            (-2, -1),
-            "backward",
-            ValueError,
-        ),
         (
             'test_n_zero',
             np.random.randn(4, 4, 4) + 1j * np.random.randn(4, 4, 4),
@@ -2009,5 +2001,245 @@ def test_grad_shape(self):
             )
 
 
+@place(DEVICES)
+@parameterize(
+    (TEST_CASE_NAME, 'x', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(3, 3, 0, 2), (1, 2), (0, 1), 'backward'),
+    ],
+)
+class TestFft2_ZeroSize(unittest.TestCase):
+    def test_fft2(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.fft2(self.x, self.n, self.axis, self.norm),
+                paddle.fft.fft2(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.fft2(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(4, 0, 6), (2, 4), None, 'backward'),
+    ],
+)
+class TestFftn_ZeroSize(unittest.TestCase):
+    def test_fftn(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.fftn(self.x, self.n, self.axis, self.norm),
+                paddle.fft.fftn(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.fftn(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(3, 3, 0, 2), (1, 2), (0, 1), 'backward'),
+    ],
+)
+class TestIfft2_ZeroSize(unittest.TestCase):
+    def test_ifft2(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.ifft2(self.x, self.n, self.axis, self.norm),
+                paddle.fft.ifft2(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.ifft2(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(4, 0, 6), (2, 4), None, 'backward'),
+    ],
+)
+class TestIfftn_ZeroSize(unittest.TestCase):
+    def test_ifftn(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.ifftn(self.x, self.n, self.axis, self.norm),
+                paddle.fft.ifftn(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.ifftn(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(3, 3, 0, 2), None, (0, 1), 'backward'),
+    ],
+)
+class TestIhfft2_ZeroSize(unittest.TestCase):
+    def test_ihfft2(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.ihfft2(self.x, self.n, self.axis, self.norm),
+                paddle.fft.ihfft2(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.ihfft2(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(4, 0, 6), (2, 4), None, 'backward'),
+    ],
+)
+class TestIhfftn_ZeroSize(unittest.TestCase):
+    def test_ihfftn(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.ihfftn(self.x, self.n, self.axis, self.norm),
+                paddle.fft.ihfftn(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.ihfftn(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(3, 3, 0, 2), None, (0, 1), 'backward'),
+    ],
+)
+class TestRfft2_ZeroSize(unittest.TestCase):
+    def test_rfft2(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.rfft2(self.x, self.n, self.axis, self.norm),
+                paddle.fft.rfft2(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.rfft2(x, self.n, self.axis, self.norm)
+            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', 'n', 'axis', 'norm'),
+    [
+        ('test_x', np.random.randn(4, 0, 6), (2, 4), None, 'backward'),
+    ],
+)
+class TestRfftn_ZeroSize(unittest.TestCase):
+    def test_rfftn(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                scipy.fft.rfftn(self.x, self.n, self.axis, self.norm),
+                paddle.fft.rfftn(
+                    paddle.to_tensor(self.x), self.n, self.axis, self.norm
+                ),
+                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.rfftn(x, self.n, self.axis, self.norm)
+            loss = paddle.sum(y)
+            loss.backward()
+            np.testing.assert_equal(
+                x.grad.shape, self.x.shape, "Grad shape mismatch"
+            )
+
+
 if __name__ == '__main__':
     unittest.main()
