```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.
#
# v4 kernel: each thread block owns ALL output channels for a set of
# frequencies, so the input spectrum is read exactly once.  512 threads =
# 16 warps, one warp per batch element; each thread accumulates a 1x4
# register tile of Y with float4 shared-memory loads of a fully cached
# W[f] slice.  The input FFT runs on the native [B, I, D, H, W] layout
# (no python-side permute/copy).
#
# 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
# ---------------------------------------------------------------------------

_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 512
#define FTILE 2

__device__ __forceinline__ float2 cmul_add(float xr, float xi,
                                           float wx, float wy, float2 acc) {
    acc.x = fmaf(xr, wx, fmaf(-xi, wy, acc.x));
    acc.y = fmaf(xr, wy, fmaf( xi, wx, acc.y));
    return acc;
}

// -------------------- 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 += 256) {
            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 += 256) {
            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[8];
    __shared__ float sh2[8];
    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 = 8;
        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, 256, 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 layout), W: [F, I, O] complex64.
// Every block owns ALL O output channels for a contiguous strided range of
// frequency tiles, so the input spectrum is read exactly once.  512 threads
// = 16 warps, warp b handles batch element b; each thread accumulates a 1x4
// register tile (o0 = 4*lane .. +3) of Y[b,o] = sum_i X[b,i,f] * W[f,i,o],
// folds the Parseval multiplicity (w-index 0 / Nyquist -> 1, else 2) and
// atomically accumulates the power at the end.  The block that owns f == 0
// also writes the DC component (spatial sums).
// Constraints checked host-side: B == 16, O == 128, I % 4 == 0, I <= 64.
__global__ void __launch_bounds__(BLOCK_SIZE)
freq_gemm_parseval_kernel(const float2* __restrict__ X,
                          const float2* __restrict__ W,
                          float* __restrict__ sums,
                          float* __restrict__ sumsq,
                          int F, int I, int O,
                          int Wh, int nyq, float inv_n) {
    extern __shared__ float2 smem[];
    const int IP = I + 2;                      // padded row (float4 aligned)
    float2* sX = smem;                         // [FTILE][16][IP]
    float2* sW = smem + FTILE * 16 * IP;       // [I][128]

    const int tid = threadIdx.x;
    const int b = tid >> 5;                    // warp id = batch index
    const int o0 = (tid & 31) << 2;            // 0..124

    const int BI = 16 * I;
    const int nF4 = (F + FTILE - 1) / FTILE;
    const int tiles = (nF4 + gridDim.x - 1) / gridDim.x;
    const int f_begin = blockIdx.x * tiles * FTILE;
    const int f_end = min(F, f_begin + tiles * FTILE * FTILE);

    float pwr0 = 0.f, pwr1 = 0.f, pwr2 = 0.f, pwr3 = 0.f;

    for (int f0 = f_begin; f0 < f_end; f0 += FTILE) {
        const int ft = min(FTILE, F - f0);
        // ---- coalesced load of FTILE frequency columns of X -------------
        for (int idx = tid; idx < BI * FTILE; idx += BLOCK_SIZE) {
            const int df = idx % FTILE;        // FTILE consecutive threads
            const int bi = idx / FTILE;        //   -> FTILE*8 contiguous bytes
            const int bb = bi / I;
            const int i = bi - bb * I;
            float2 v = make_float2(0.f, 0.f);
            if (df < ft) v = X[(long)bi * F + (f0 + df)];
            sX[(df * 16 + bb) * IP + i] = v;
        }
        __syncthreads();

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

            // ---- coalesced load of W[f] (I x 128 complex) --------------
            {
                const float4* W4 = reinterpret_cast<const float4*>(
                    W + (long)f * I * O);
                float4* sW4 = reinterpret_cast<float4*>(sW);
                const int n4 = (I * O) >> 1;
                for (int t = tid; t < n4; t += BLOCK_SIZE) sW4[t] = W4[t];
            }
            __syncthreads();

            const float2* xr = sX + (df * 16 + b) * IP;
            float2 acc0 = make_float2(0.f, 0.f);
            float2 acc1 = make_float2(0.f, 0.f);
            float2 acc2 = make_float2(0.f, 0.f);
            float2 acc3 = make_float2(0.f, 0.f);

            #pragma unroll 4
            for (int i2 = 0; i2 < (I >> 1); ++i2) {
                const float4 xa = *reinterpret_cast<const float4*>(
                    xr + (i2 << 1));
                const float2* w0p = sW + (((i2 << 1)    ) << 7) + o0;
                const float2* w1p = sW + (((i2 << 1) + 1) << 7) + o0;
                const float4 wa = *reinterpret_cast<const float4*>(w0p);
                const float4 wc = *reinterpret_cast<const float4*>(w0p + 2);
                const float4 wd = *reinterpret_cast<const float4*>(w1p);
                const float4 we = *reinterpret_cast<const float4*>(w1p + 2);
                acc0 = cmul_add(xa.x, xa.y, wa.x, wa.y, acc0);
                acc1 = cmul_add(xa.x, xa.y, wa.z, wa.w, acc1);
                acc2 = cmul_add(xa.x, xa.y, wc.x, wc.y, acc2);
                acc3 = cmul_add(xa.x, xa.y, wc.z, wc.w, acc3);
                acc0 = cmul_add(xa.z, xa.w, wd.x, wd.y, acc0);
                acc1 = cmul_add(xa.z, xa.w, wd.z, wd.w, acc1);
                acc2 = cmul_add(xa.z, xa.w, we.x, we.y, acc2);
                acc3 = cmul_add(xa.z, xa.w, we.z, we.w, acc3);
            }

            if (f == 0) {                      // DC component = spatial sum
                float* srow = sums + b * O + o0;
                srow[0] = acc0.x;
                srow[1] = acc1.x;
                srow[2] = acc2.x;
                srow[3] = acc3.x;
            }
            pwr0 += mult * (acc0.x * acc0.x + acc0.y * acc0.y);
            pwr1 += mult * (acc1.x * acc1.x + acc1.y * acc1.y);
            pwr2 += mult * (acc2.x * acc2.x + acc2.y * acc2.y);
            pwr3 += mult * (acc3.x * acc3.x + acc3.y * acc3.y);
            __syncthreads();                   // protect sW
        }
        __syncthreads();                       // protect sX
    }

    float* out = sumsq + b * O + o0;
    atomicAdd(out + 0, pwr0);
    atomicAdd(out + 1, pwr1);
    atomicAdd(out + 2, pwr2);
    atomicAdd(out + 3, pwr3);
}

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 == 16 && O == 128 && (I % 4) == 0 && I <= 64,
                "unsupported shape for fused kernel");
    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};

    const int IP = (int)I + 2;
    const int smem = (FTILE * 16 * IP + (int)(I * O)) * 8;
    cudaError_t err = cudaFuncSetAttribute(
        freq_gemm_parseval_kernel,
        cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
    TORCH_CHECK(err == cudaSuccess,
                "insufficient shared memory for fused kernel: ",
                cudaGetErrorString(err));
    auto prop = at::cuda::getCurrentDeviceProperties();
    int nF4 = (int)((F + FTILE - 1) / FTILE);
    int grid = std::min(nF4, prop->multiProcessorCount * 2);
    auto stream = at::cuda::getCurrentCUDAStream();
    freq_gemm_parseval_kernel<<<grid, BLOCK_SIZE, smem, 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)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_v4",
    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.

    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

    def _fft_shared_ok(self, x, In):
        # dynamic shared memory needed by the fused kernel
        need = (_F_TILE * 16 * (In + 2) + In * 128) * 8
        try:
            props = torch.cuda.get_device_properties(x.device)
            limit = getattr(props, "shared_memory_per_block_optin", None)
            if limit is None:
                limit = getattr(props, "shared_memory_per_block", 0)
            return need <= int(limit)
        except Exception:
            return True  # let the extension try (it checks and raises)

    # ---------------------------------------------------------------- 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 B == 16 and Out == 128
                and In % 4 == 0 and 0 < In <= 64
                and self._fft_shared_ok(x, In)):
            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

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