fused_fft_conv2d#

fused_fft_conv2d(x, weight, fft_size=None)#

Computes 2D depthwise convolution via the fused FFT pipeline — with a shared filter or one filter per batch element — with autograd support for both inputs: the filter spectrum kf = fused_rfft2d_fwd(weight, fft_size) is computed once per call, then y = fused_fft_conv2d_fwd(x, kf).

The filter is treated as a full (fft_size, fft_size) tile (zero-padded bottom/right) and the output is the ‘same’-centered crop of the FFT convolution at offset fft_size // 2 — equivalent to fft_conv2d() with the weight zero-padded to (fft_size, fft_size):

w_fs = F.pad(weight, (0, fft_size - K_y, 0, fft_size - K_x))
y = torch.fft.irfft2(
    torch.fft.rfft2(x, s=(fft_size, fft_size))
    * torch.fft.rfft2(w_fs, s=(fft_size, fft_size)),
    s=(fft_size, fft_size),
)[..., fft_size // 2 :, fft_size // 2 :][..., : x.shape[-2], : x.shape[-1]]
Parameters:

  • x (torch.Tensor) – Input tensor of shape (batch_size, hidden_dim, x_dim_seq, y_dim_seq) with the sequence dims at most fft_size // 2.

  • weight (torch.Tensor) – Weight tensor of shape (1, hidden_dim, x_dim_kernel, y_dim_kernel) (shared filter), (batch_size, hidden_dim, x_dim_kernel, y_dim_kernel) (per-batch filters), or 3D (hidden_dim, x_dim_kernel, y_dim_kernel) (treated as shared), with the kernel dims at most fft_size. Must have the same dtype as x.

  • fft_size (int, optional) – Square FFT tile size; must be a power of two in {8, 16, 32, 64, 128}. Defaults to 2 * next_pow2(max(x_dim_seq, y_dim_seq)).

Returns:

Output tensor of shape (batch_size, hidden_dim, x_dim_seq, y_dim_seq) in the same dtype as x. Differentiable with respect to x and weight (the backward runs the fused dx+dkf kernel, then maps dkf to the spatial filter gradient with fused_rfft2d_bwd()).

Return type:

torch.Tensor

On this page