Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wanghuancoder merged 5 commits into from
Jun 4, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9f84878
Fix
co63oc Jun 1, 2025
7c8caee
Fix
co63oc Jun 1, 2025
b43e8d5
Fix
co63oc Jun 1, 2025
560e75d
Fix
co63oc Jun 3, 2025
f05351a
Fix
co63oc Jun 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions paddle/phi/infermeta/unary.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]));
Expand All @@ -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]]));
Expand All @@ -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));
Expand Down Expand Up @@ -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]));
Expand Down
11 changes: 11 additions & 0 deletions paddle/phi/kernels/impl/fft_grad_kernel_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;

Expand Down Expand Up @@ -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;
Expand Down
24 changes: 23 additions & 1 deletion paddle/phi/kernels/impl/fft_kernel_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -33,6 +33,18 @@ void FFTC2CKernel(const Context& ctx,
bool forward,
DenseTensor* out) {
ctx.template Alloc<T>(out);
if (x.numel() == 0) {

@DanielSun11 DanielSun11 Jun 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

涉及到这几个fft相关的kernel是不是input和output的shape是相同的？若是，x.numel() == 0时，out也为0 size Tensor，这时直接return就行吧？此处调用Full 是另有什么考虑吗？

@co63oc co63oc Jun 3, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

>>> 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]]]])

这样执行时会返回为0，这里不容易理解，已增加到注释

@DanielSun11 DanielSun11 Jun 4, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok,

void FFTR2CInferMeta(const MetaTensor& x,
                     const std::vector<int64_t>& axes,
                     const std::string& normalization,
                     bool forward,
                     bool onesided,
                     MetaTensor* out,
                     MetaConfig config) {
  PADDLE_ENFORCE_NOT_NULL(
      out,
      common::errors::InvalidArgument("Output of fft_r2c should not be null."));
  const phi::DDim x_dim = x.dims();

  // only ensure that fft axes' size greater than zero at runtime
  // 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],
                        0,
                        common::errors::InvalidArgument(
                            "Invalid fft n-point (%d).", x_dim[axis]));
    }
  }

  out->set_layout(x.layout());
  out->set_dtype(ToComplexType(x.dtype()));
  if (!onesided) {
    out->share_dims(x);
  } else {
    phi::DDim out_dim = x.dims();
    const int last_fft_axis = static_cast<int>(axes.back());
    const int64_t last_fft_dim_size = x_dim[last_fft_axis];
    out_dim.at(last_fft_axis) = last_fft_dim_size / 2 + 1;
    out->set_dims(out_dim);
  }
}

fft相关的kernel input和output的shape不一定相同

/*
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);
Expand All @@ -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);
Expand All @@ -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;

Expand Down
248 changes: 240 additions & 8 deletions test/fft/test_fft.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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()