```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Algorithmic rewrite of: ConvTranspose3d -> scale -> BatchNorm3d -> global
# average pooling.
#
# The network ends in a *global* average pool and BatchNorm is a per-channel
# affine transform, so the result depends only on the per-(batch, channel)
# spatial sums and sums-of-squares of the transposed-convolution output.
#
# For stride-1 / pad-0 ConvTranspose3d the output is the *full* linear
# convolution of x and w.  Hence (convolution theorem + Parseval):
#     sum_z y[c,z]   = Y[c,0]                        (DC Fourier component)
#     sum_z y[c,z]^2 = (1/N) * sum_f m(f) * |Y[c,f]|^2
# with Y[c,f] = sum_i X[i,f] * W[i,c,f] -- a small complex GEMM per frequency.
# The huge conv output (~200 MB) and the ~850 GFLOP direct convolution are
# never materialized; the pointwise |Y|^2 reduction is fused into the GEMM
# kernel, and a tiny finalize kernel applies conv bias, scale, BatchNorm
# (batch stats + running-stat update, or running stats) and global pooling.
#
# Eval mode (fixed running stats) uses an even cheaper exact closed form:
#     sum(out[b,o]) = sum_i sum(x[b,i]) * sum(w[i,o]) + bias[o] * S_out
#
# This revision fuses the [F,B,I] spectral transpose INTO the fused kernel:
# cuFFT is run on the native [B,I,D,H,W] layout (no python-side permute /
# copy) and each block stages a 2-frequency tile of X into shared memory
# with fully coalesced reads.  The weight spectrum W[F,I,O] is read exactly
# once (a single block row covers all frequencies).
# ---------------------------------------------------------------------------

_F_TILE = 2  # frequencies staged per shared-memory tile (keep in sync w/ CUDA)

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/util/complex.h>
#include <cuda_runtime.h>
#include <vector>
#include <algorithm>

#define BLOCK_SIZE 256
#define MAX_PAIRS 8
#define MAX_BI 2048
#define MAX_I 64
#define FTILE 2

// -------------------- per-(b,c) spatial sums / sumsq ----------------------
__global__ void spatial_sums_kernel(const float* __restrict__ x,
                                    float* __restrict__ sums,
                                    float* __restrict__ sumsq,
                                    int S, int compute_sq) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    long base = ((long)b * (long)gridDim.x + (long)c) * (long)S;
    const float* xp = x + base;

    float acc = 0.f, acc2 = 0.f;
    if ((S & 3) == 0) {
        const float4* xp4 = reinterpret_cast<const float4*>(xp);
        int S4 = S >> 2;
        for (int s = threadIdx.x; s < S4; s += BLOCK_SIZE) {
            float4 v = xp4[s];
            acc += (v.x + v.y) + (v.z + v.w);
            if (compute_sq) acc2 += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
        }
    } else {
        for (int s = threadIdx.x; s < S; s += BLOCK_SIZE) {
            float v = xp[s];
            acc += v;
            if (compute_sq) acc2 += v * v;
        }
    }

    for (int off = 16; off > 0; off >>= 1) {
        acc += __shfl_down_sync(0xffffffffu, acc, off);
        if (compute_sq) acc2 += __shfl_down_sync(0xffffffffu, acc2, off);
    }
    __shared__ float sh[BLOCK_SIZE / 32];
    __shared__ float sh2[BLOCK_SIZE / 32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh[wid] = acc;
        if (compute_sq) sh2[wid] = acc2;
    }
    __syncthreads();
    if (wid == 0) {
        int nw = BLOCK_SIZE / 32;
        float v  = (lane < nw) ? sh[lane] : 0.f;
        float v2 = (lane < nw && compute_sq) ? sh2[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
            if (compute_sq) v2 += __shfl_down_sync(0xffffffffu, v2, off);
        }
        if (lane == 0) {
            long idx = (long)b * (long)gridDim.x + (long)c;
            sums[idx] = v;
            if (compute_sq) sumsq[idx] = v2;
        }
    }
}

std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    int64_t B = xc.size(0);
    int64_t C = xc.size(1);
    int64_t S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    TORCH_CHECK(B <= 65535, "batch too large");
    auto sums  = torch::empty({B, C}, xc.options());
    auto sumsq = compute_sq ? torch::empty({B, C}, xc.options())
                            : torch::empty({0}, xc.options());
    if (B * C == 0) return {sums, sumsq};
    dim3 grid((unsigned)C, (unsigned)B);
    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sums_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), sums.data_ptr<float>(),
        compute_sq ? sumsq.data_ptr<float>() : nullptr,
        (int)S, compute_sq ? 1 : 0);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {sums, sumsq};
}

// ------------- fused frequency GEMM + Parseval power reduction ------------
// X: [B, I, F] complex64 (native rfftn output layout -- no transpose copy),
// W: [F, I, O] complex64 (cached weight spectrum).
// Block (fx, oy) handles frequency *tiles* (FTILE consecutive frequencies)
// t = fx, fx+gx, ... and output-channel chunk oc = 32*oy.  For each tile it
// stages the FTILE x (B*I) slice of X into shared memory with fully
// coalesced global reads (FTILE consecutive complex elements per row), then
// computes Y[b,o] = sum_i X[b,i,f]*W[f,i,o] for all b and the 32 o's of the
// chunk, folds the Parseval multiplicity (w-index 0 / Nyquist -> 1, else 2),
// accumulates per-thread partials in registers across its tiles, and finally
// atomically adds them into sumsq.  The block row that owns f == 0 also
// writes the DC component (spatial sums).  Since oy == 0 only, W is read
// exactly once.  Constraints (checked host-side): B*I <= MAX_BI, I <= MAX_I,
// O == 32, B <= 8*BLOCK_SIZE/32.
__global__ void freq_gemm_parseval_kernel(const float2* __restrict__ X,
                                          const float2* __restrict__ W,
                                          float* __restrict__ sums,
                                          float* __restrict__ sumsq,
                                          int F, int B, int I, int O,
                                          int Wh, int nyq, float inv_n) {
    __shared__ float2 sX[FTILE * MAX_BI];
    __shared__ float2 sW[MAX_I * 32];

    const int oc = blockIdx.y << 5;
    const int BI = B * I;

    // fixed (b, o_local) pairs handled by this thread
    int pb[MAX_PAIRS], po[MAX_PAIRS];
    bool pv[MAX_PAIRS];
    #pragma unroll
    for (int k = 0; k < MAX_PAIRS; ++k) {
        int p = k * BLOCK_SIZE + threadIdx.x;
        pv[k] = p < (B << 5);
        pb[k] = pv[k] ? (p >> 5) : 0;
        po[k] = p & 31;
    }

    float pwr[MAX_PAIRS];
    #pragma unroll
    for (int k = 0; k < MAX_PAIRS; ++k) pwr[k] = 0.f;

    const int nTiles = (F + FTILE - 1) / FTILE;

    for (int t0 = blockIdx.x; t0 < nTiles; t0 += gridDim.x) {
        const int f0 = t0 * FTILE;
        const int ft = min(FTILE, F - f0);

        // ---- stage FTILE frequency columns of X (coalesced) --------------
        for (int idx = threadIdx.x; idx < BI * FTILE; idx += BLOCK_SIZE) {
            const int df = idx % FTILE;    // FTILE consecutive threads read
            const int bi = idx / FTILE;    // FTILE*8 contiguous bytes
            float2 v = make_float2(0.f, 0.f);
            if (df < ft) v = X[(long)bi * F + (f0 + df)];
            sX[df * BI + bi] = v;
        }

        // ---- stage W[f] chunks (coalesced 32-wide rows) -------------------
        const int IW = I << 5;
        for (int t = threadIdx.x; t < FTILE * IW; t += BLOCK_SIZE) {
            const int df = t / IW;
            const int r  = t - df * IW;
            float2 v = make_float2(0.f, 0.f);
            if (df < ft) {
                const long fbaseW = (long)(f0 + df) * I * O + oc;
                v = W[fbaseW + (long)(r >> 5) * O + (r & 31)];
            }
            sW[t] = v;
        }

        __syncthreads();

        #pragma unroll
        for (int df = 0; df < FTILE; ++df) {
            const int f = f0 + df;
            if (df >= ft) break;
            const int w = f - (f / Wh) * Wh;
            const float mult = ((w == 0) || (w == nyq)) ? inv_n : 2.0f * inv_n;

            #pragma unroll
            for (int k = 0; k < MAX_PAIRS; ++k) {
                if (pv[k]) {
                    float2 a = make_float2(0.f, 0.f);
                    const float2* xrow = sX + df * BI + pb[k] * I;
                    const float2* wcol = sW + df * IW + po[k];
                    #pragma unroll 4
                    for (int i = 0; i < I; ++i) {
                        float2 xv = xrow[i];
                        float2 wv = wcol[i << 5];
                        a.x += xv.x * wv.x - xv.y * wv.y;
                        a.y += xv.x * wv.y + xv.y * wv.x;
                    }
                    if (f == 0) sums[pb[k] * O + oc + po[k]] = a.x;  // DC
                    pwr[k] += mult * (a.x * a.x + a.y * a.y);
                }
            }
        }
        __syncthreads();
    }

    #pragma unroll
    for (int k = 0; k < MAX_PAIRS; ++k) {
        if (pv[k]) atomicAdd(&sumsq[pb[k] * O + oc + po[k]], pwr[k]);
    }
}

std::vector<torch::Tensor> freq_gemm_parseval(torch::Tensor X, torch::Tensor W,
                                              int64_t B, int64_t I, int64_t O,
                                              int64_t Wh, int64_t nyq,
                                              double inv_n) {
    TORCH_CHECK(X.scalar_type() == at::kComplexFloat, "X must be complex64");
    TORCH_CHECK(W.scalar_type() == at::kComplexFloat, "W must be complex64");
    TORCH_CHECK(X.is_contiguous() && W.is_contiguous(), "X, W must be contiguous");
    TORCH_CHECK(B * I <= MAX_BI && I <= MAX_I && (O % 32) == 0 &&
                B <= (MAX_PAIRS * BLOCK_SIZE) / 32, "unsupported shape");
    TORCH_CHECK(X.dim() == 3, "X must be [B, I, F]");
    int64_t F = X.size(2);
    auto opts = X.options().dtype(at::kFloat);
    auto sums  = torch::empty({B, O}, opts);
    auto sumsq = torch::zeros({B, O}, opts);
    if (F == 0 || B * O == 0) return {sums, sumsq};

    int sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    int64_t nTiles = (F + FTILE - 1) / FTILE;
    int f_blocks = (int)std::min<int64_t>(nTiles, (int64_t)sm * 8);
    dim3 grid((unsigned)f_blocks, (unsigned)(O / 32));
    auto stream = at::cuda::getCurrentCUDAStream();
    freq_gemm_parseval_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        reinterpret_cast<const float2*>(X.data_ptr<c10::complex<float>>()),
        reinterpret_cast<const float2*>(W.data_ptr<c10::complex<float>>()),
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        (int)F, (int)B, (int)I, (int)O, (int)Wh, (int)nyq, (float)inv_n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {sums, sumsq};
}

// ------------------------------ finalize ----------------------------------
// Given raw per-(b,c) spatial sums and sums-of-squares of the *unbiased*
// conv output, applies the conv-bias correction, the scale, BatchNorm (batch
// statistics + running-stat update exactly like PyTorch when training, else
// running statistics) and emits the globally pooled result.
__global__ void finalize_kernel(const float* __restrict__ sum0,
                                const float* __restrict__ sumsq0,
                                const float* __restrict__ bias,
                                float* __restrict__ rmean,
                                float* __restrict__ rvar,
                                const float* __restrict__ gamma,
                                const float* __restrict__ beta,
                                float* __restrict__ out,
                                int B, int C, float S,
                                float scale, float momentum,
                                float eps, int training) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * C) return;
    int b = idx / C;
    int c = idx - b * C;
    float bc = (bias != nullptr) ? bias[c] : 0.f;
    float s_raw  = sum0[idx];
    float s2_raw = sumsq0[idx];
    float su  = s_raw + bc * S;
    float s2u = s2_raw + 2.f * bc * s_raw + bc * bc * S;
    float mean, invstd;
    float N = S * (float)B;
    if (training) {
        float s = 0.f, s2 = 0.f;
        for (int bb = 0; bb < B; ++bb) {
            float sr  = sum0[(long)bb * C + c];
            float s2r = sumsq0[(long)bb * C + c];
            s  += sr + bc * S;
            s2 += s2r + 2.f * bc * sr + bc * bc * S;
        }
        s  *= scale;
        s2 *= scale * scale;
        mean = s / N;
        float var = s2 / N - mean * mean;
        var = fmaxf(var, 0.f);
        invstd = rsqrtf(var + eps);
        if (b == 0 && rmean != nullptr && momentum >= 0.f) {
            float unbiased = var * (N / (N - 1.f));
            rmean[c] = (1.f - momentum) * rmean[c] + momentum * mean;
            rvar[c]  = (1.f - momentum) * rvar[c]  + momentum * unbiased;
        }
    } else {
        mean = rmean[c];
        invstd = rsqrtf(rvar[c] + eps);
    }
    float g  = (gamma != nullptr) ? gamma[c] : 1.f;
    float be = (beta  != nullptr) ? beta[c]  : 0.f;
    float mbc = su * (scale / S);
    out[idx] = (mbc - mean) * invstd * g + be;
}

torch::Tensor finalize(torch::Tensor sums, torch::Tensor sumsq,
                       torch::Tensor bias, torch::Tensor rmean,
                       torch::Tensor rvar, torch::Tensor gamma,
                       torch::Tensor beta, double scale, double S,
                       double momentum, double eps, bool training) {
    int B = sums.size(0), C = sums.size(1);
    auto out = torch::empty({B, C}, sums.options());
    int total = B * C;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.numel() ? rmean.data_ptr<float>() : nullptr,
        rvar.numel()  ? rvar.data_ptr<float>()  : nullptr,
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, C, (float)S, (float)scale,
        (float)momentum, (float)eps, training ? 1 : 0);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

// Eval closed form: per-channel sums of the conv output via
// sum(out[b,o]) = sum_i Sx[b,i] * Sw[i,o], then +bias, scale, BN, pool.
__global__ void eval_finalize_kernel(const float* __restrict__ Sx,
                                     const float* __restrict__ Sw,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ rmean,
                                     const float* __restrict__ rvar,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     float* __restrict__ out,
                                     int B, int In, int Out,
                                     float scale, float inv_Sout, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * Out) return;
    int b = idx / Out;
    int o = idx - b * Out;
    const float* sx = Sx + (long)b * In;
    float acc = 0.f;
    for (int i = 0; i < In; ++i) acc += sx[i] * Sw[(long)i * Out + o];
    float mean = acc * inv_Sout;
    if (bias != nullptr) mean += bias[o];
    mean *= scale;
    float g  = (gamma != nullptr) ? gamma[o] : 1.f;
    float be = (beta  != nullptr) ? beta[o]  : 0.f;
    out[idx] = (mean - rmean[o]) * rsqrtf(rvar[o] + eps) * g + be;
}

torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta, double scale,
                            double inv_Sout, double eps) {
    int B = Sx.size(0), In = Sx.size(1), Out = Sw.size(1);
    auto out = torch::empty({B, Out}, Sx.options());
    int total = B * Out;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    eval_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, In, Out,
        (float)scale, (float)inv_Sout, (float)eps);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_cpp_src = r"""
#include <vector>
std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq);
std::vector<torch::Tensor> freq_gemm_parseval(torch::Tensor X, torch::Tensor W,
                                              int64_t B, int64_t I, int64_t O,
                                              int64_t Wh, int64_t nyq,
                                              double inv_n);
torch::Tensor finalize(torch::Tensor sums, torch::Tensor sumsq,
                       torch::Tensor bias, torch::Tensor rmean,
                       torch::Tensor rvar, torch::Tensor gamma,
                       torch::Tensor beta, double scale, double S,
                       double momentum, double eps, bool training);
torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta, double scale,
                            double inv_Sout, double eps);
"""

_ext = load_inline(
    name="fused_fft_convtranspose3d_bn_globalpool_v5",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["spatial_sums", "freq_gemm_parseval", "finalize", "eval_finalize"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Drop-in optimized replacement for Model (identical submodules/state dict).

    Train mode (or eval without tracked stats): the ConvTranspose3d is
    evaluated in the Fourier domain (a stride-1/pad-0 transposed convolution
    is a full linear convolution == pointwise product of padded spectra).
    Only the DC component (spatial sums) and the Parseval power (sums of
    squares) of the output spectrum are computed -- fused into one custom
    CUDA kernel -- and a tiny finalize kernel applies conv bias, scale,
    BatchNorm (with exact running-stat update) and global average pooling.
    The ~200 MB conv output and the ~850 GFLOP direct convolution are never
    materialized.  The input FFT runs on the native [B,I,D,H,W] layout and
    the spectral transpose is fused into the kernel (no permute copies).

    Eval mode with tracked running stats: spatial sums of the conv output
    are computed in exact closed form from per-channel input sums and kernel
    sums, so no convolution runs at all.
    """

    def __init__(self, in_channels, out_channels, kernel_size, scale_factor,
                 eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = float(scale_factor)
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        ct = self.conv_transpose
        self._simple_conv = (
            ct.groups == 1
            and tuple(ct.stride) == (1, 1, 1)
            and tuple(ct.padding) == (0, 0, 0)
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
        )
        self._w_spec = None   # cached weight spectrum [F, In, Out] complex64
        self._w_key = None

    # ------------------------------------------------------------------ util
    def _reference_forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

    def _out_spatial(self, x):
        k = self.conv_transpose.kernel_size
        return (x.shape[2] + k[0] - 1, x.shape[3] + k[1] - 1, x.shape[4] + k[2] - 1)

    def _get_weight_spectrum(self, s):
        w = self.conv_transpose.weight
        key = (w._version, w.data_ptr(), s)
        if self._w_key != key or self._w_spec is None:
            wp = w.detach().permute(2, 3, 4, 0, 1)                 # [kd,kh,kw,In,Out]
            spec = torch.fft.rfftn(wp, s=s, dim=(0, 1, 2))         # [D,H,Wh,In,Out]
            F = spec.shape[0] * spec.shape[1] * spec.shape[2]
            self._w_spec = spec.reshape(F, spec.shape[3], spec.shape[4]).contiguous()
            self._w_key = key
        return self._w_spec

    # ---------------------------------------------------------------- forward
    def forward(self, x):
        bn = self.batch_norm
        ct = self.conv_transpose

        if not (x.is_cuda and x.dtype == torch.float32):
            return self._reference_forward(x)
        if bn.training and bn.momentum is None:
            return self._reference_forward(x)  # cumulative moving average
        if (not bn.training) and bn.track_running_stats and (
            bn.running_mean is None or not bn.running_mean.is_cuda
        ):
            return self._reference_forward(x)

        B = x.shape[0]
        In, Out = ct.in_channels, ct.out_channels
        gamma = bn.weight if bn.affine else x.new_empty(0)
        beta = bn.bias if bn.affine else x.new_empty(0)
        bias = ct.bias if ct.bias is not None else x.new_empty(0)
        if bn.track_running_stats:
            rmean, rvar = bn.running_mean, bn.running_var
        else:
            rmean = rvar = x.new_empty(0)
        momentum = bn.momentum if bn.momentum is not None else -1.0
        use_batch_stats = bn.training or (not bn.track_running_stats)

        if (not bn.training) and bn.track_running_stats and self._simple_conv:
            # ---- closed-form eval path: no convolution at all --------------
            Sx = _ext.spatial_sums(x, False)[0]            # [B, In]
            Sw = _ext.spatial_sums(ct.weight, False)[0]    # [In, Out]
            Do, Ho, Wo = self._out_spatial(x)
            S_out = Do * Ho * Wo
            out = _ext.eval_finalize(
                Sx, Sw, bias, rmean, rvar, gamma, beta,
                self.scale_factor, 1.0 / float(S_out), bn.eps,
            )
            return out.view(B, Out, 1, 1, 1)

        # ------------------------ FFT path (train / no-track eval) -----------
        if self._simple_conv and (In <= 64 and Out % 32 == 0
                                  and B <= 64 and B * In <= 2048):
            Do, Ho, Wo = self._out_spatial(x)
            s = (Do, Ho, Wo)
            Wh = Wo // 2 + 1
            F = Do * Ho * Wh
            N = Do * Ho * Wo
            nyq = (Wh - 1) if (Wo % 2 == 0) else -1

            # Native-layout FFT: no permute / transpose copies anywhere.
            Xn = torch.fft.rfftn(x, s=s, dim=(2, 3, 4))    # [B,In,Do,Ho,Wh]
            Xn = Xn.reshape(B, In, F)                      # contiguous view
            Wf = self._get_weight_spectrum(s)              # [F, In, Out]

            sums, sumsq = _ext.freq_gemm_parseval(
                Xn, Wf, B, In, Out, Wh, nyq, 1.0 / float(N))

            if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
                bn.num_batches_tracked += 1
            out = _ext.finalize(
                sums, sumsq, bias, rmean, rvar, gamma, beta,
                self.scale_factor, float(N), momentum, bn.eps, use_batch_stats)
            return out.view(B, Out, 1, 1, 1)

        # ---- generic fallback: run the conv, fuse the rest ------------------
        y = ct(x)
        C = y.shape[1]
        S = y.shape[2] * y.shape[3] * y.shape[4]
        sums, sumsq = _ext.spatial_sums(y, True)
        if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
            bn.num_batches_tracked += 1
        out = _ext.finalize(
            sums, sumsq, x.new_empty(0), rmean, rvar, gamma, beta,
            self.scale_factor, float(S), momentum, bn.eps, use_batch_stats)
        return out.view(B, C, 1, 1, 1)
```