mixed_fftconv2d_fp32_bhl_w_reshape#

mixed_fftconv2d_fp32_bhl_w_reshape(
x,
kernel,
periodic,
shortcut=None,
use_phase_shift=True,
)#

2D mixed-BC FFT conv wrapper for BLH layout (batch, X, Y, hidden).

Reshapes BLH ↔ BHL around mixed_fftconv2d_fp32_bhl(). Prefer this over operating in BLH natively — the FFT runs faster on contiguous spatial axes (BHL).

Parameters:

  • x (Tensor) – Input tensor of shape [B, X, Y, H] (BLH, channels-last).

  • kernel (Tensor) – Kernel tensor of shape [1|B, K_x, K_y, H] (BLH). Leading dim 1 for a shared kernel, B for per-sample kernels.

  • periodic (Sequence[bool]) – Length-2 sequence (periodic_x, periodic_y).

  • shortcut (Tensor | None) – Optional per-channel scale [H].

  • use_phase_shift (bool) – See mixed_fftconv1d_fp32_bhl().

Returns:

Tensor of shape [B, X, Y, H] in the original dtype of x.

Return type:

Tensor