You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # 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
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # 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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

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

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, scales the output, applies batch normalization, 
    and then performs global average pooling. 
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def 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

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 16, 32, 32
kernel_size = 5
scale_factor = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scale_factor]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 5 of 6 for this task. Your previous solution (turn 4) was:
```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
#
# v3: the batched complex GEMM kernel is reworked -- each thread block owns
# all output channels for a strided set of frequencies (input spectrum is
# read once instead of 4x), uses a 2x4 register micro-tile with float4
# shared-memory loads on padded/bank-conflict-free layouts, and the input
# FFT is taken on the native [B, I, D, H, W] layout (no python-side permute
# copies) with coalesced 4-frequency-tiled global loads.
# ---------------------------------------------------------------------------

_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 4

// -------------------- 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);
    return {sums, sumsq};
}

// ------------- fused frequency GEMM + Parseval power reduction ------------
// X: [B, I, F] complex64 (native rfftn layout), W: [F, I, O] complex64.
// Each block owns ALL O output channels and a strided set of frequencies,
// loading the input spectrum once in coalesced FTILE-frequency tiles.
// Per frequency it computes Y[b,o] = sum_i X[b,i,f] * W[i,o,f] for a
// 2(b) x 4(o) register micro-tile per thread, folds the Parseval
// multiplicity (w-index 0 / Nyquist -> 1, other w -> 2) and atomically
// accumulates the power.  The block that owns f == 0 also writes the DC
// component (spatial sums).  Constraints checked host-side:
//   B == 16, I % 4 == 0, I <= 66, O == 128.
__global__ void __launch_bounds__(BLOCK_SIZE, 2)
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/2][32][2] chunk

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int lb = (lane >> 3) & 1;           // b-quad index within warp
    const int lo = lane & 7;                  // o-oct index within lane quad
    const int wb = warp >> 2;                 // 0..3 : b-quad group
    const int wo = warp & 3;                  // 0..3 : 32-o group
    const int b0 = (wb << 2) + (lb << 1);     // thread's first b
    const int o0 = (wo << 5) + lo;            // thread's o base (+8*j)

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

    float pwr[2][4];
    #pragma unroll
    for (int u = 0; u < 2; ++u)
        #pragma unroll
        for (int j = 0; j < 4; ++j) pwr[u][j] = 0.f;

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

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

            const long fbase = (long)f * I * O;
            for (int oc = 0; oc < O; oc += 32) {
                for (int t = tid; t < (I >> 1) * 32; t += BLOCK_SIZE) {
                    const int ii = t >> 5;
                    const int o = t & 31;
                    sW[(ii * 32 + o) << 1]       = W[fbase + (long)((ii << 1)    ) * O + oc + o];
                    sW[((ii * 32 + o) << 1) + 1] = W[fbase + (long)((ii << 1) + 1) * O + oc + o];
                }
                __syncthreads();

                const int ob = (oc & 96) + lo;  // oc==64 -> 64+lo (warp folded)
                float2 ar[2][4], ai[2][4];
                #pragma unroll
                for (int u = 0; u < 2; ++u)
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        ar[u][j] = make_float2(0.f, 0.f);
                        ai[u][j] = make_float2(0.f, 0.f);
                    }

                const float2* xrow0 = sX + (df * 16 + b0    ) * IP;
                const float2* xrow1 = sX + (df * 16 + b0 + 1) * IP;

                #pragma unroll 4
                for (int i2 = 0; i2 < (I >> 1); ++i2) {
                    const float4 xa = *reinterpret_cast<const float4*>(xrow0 + (i2 << 1));
                    const float4 xb = *reinterpret_cast<const float4*>(xrow1 + (i2 << 1));
                    const float2* wp = sW + (i2 << 6);
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        const float4 wv = *reinterpret_cast<const float4*>(
                            wp + ((ob + (j << 3)) << 1));
                        const float xr0 = xa.x, xi0 = xa.y, xr1 = xa.z, xi1 = xa.w;
                        const float xs0 = xb.x, xt0 = xb.y, xs1 = xb.z, xt1 = xb.w;
                        const float wr0 = wv.x, wi0 = wv.y, wr1 = wv.z, wi1 = wv.w;
                        ar[0][j].x += xr0 * wr0 - xi0 * wi0;
                        ar[0][j].y += xr0 * wi0 + xi0 * wr0;
                        ai[0][j].x += xr1 * wr1 - xi1 * wi1;
                        ai[0][j].y += xr1 * wi1 + xi1 * wr1;
                        ar[1][j].x += xs0 * wr0 - xt0 * wi0;
                        ar[1][j].y += xs0 * wi0 + xt0 * wr0;
                        ai[1][j].x += xs1 * wr1 - xt1 * wi1;
                        ai[1][j].y += xs1 * wi1 + xt1 * wr1;
                    }
                }

                #pragma unroll
                for (int u = 0; u < 2; ++u) {
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        const int b = b0 + u;
                        const int o = ob + (j << 3);
                        float2 a;
                        a.x = ar[u][j].x + ai[u][j].x;
                        a.y = ar[u][j].y + ai[u][j].y;
                        if (f == 0) sums[b * O + o] = a.x;   // DC = spatial sum
                        pwr[u][j] += mult * (a.x * a.x + a.y * a.y);
                    }
                }
                __syncthreads();
            }
        }
        __syncthreads();
    }

    #pragma unroll
    for (int u = 0; u < 2; ++u) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int b = b0 + u;
            const int o0w = (wo << 5) + lo + (j << 3);
            atomicAdd(&sumsq[b * O + o0w],      pwr[u][j]);
            atomicAdd(&sumsq[b * O + o0w + 64], pwr[u][j]);
        }
    }
}

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 <= 66,
                "unsupported shape for fused kernel");
    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 * 32)) * 8;
    static int smem_limit = -1;
    if (smem_limit < 0) {
        cudaFuncSetAttribute(freq_gemm_parseval_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             100 * 1024);
        smem_limit = 1;
    }
    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);
    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);
    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);
    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_v3",
    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

    # ---------------------------------------------------------------- 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 In % 4 == 0 and In <= 66
                                  and Out == 128):
            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)
```
Result of evaluating it:
The solution compiled but failed at run time. Error output:
```
Dimension out of range (expected to be in range of [-2, 1], but got 2)
```

The fastest correct solution so far (3.988x speedup) is:
```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
# ---------------------------------------------------------------------------

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

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

// -------------------- 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);
    return {sums, sumsq};
}

// ------------- fused frequency GEMM + Parseval power reduction ------------
// X: [F, B, I] complex64, W: [F, I, O] complex64.
// Block (fx, oy) handles frequencies f = fx, fx+gx, ... and output-channel
// chunk oc = 32*oy.  For each f it computes Y[b,o] = sum_i X[f,b,i]*W[f,i,o]
// for all b and the 32 o's of the chunk, folds the Parseval multiplicity
// (w-index 0 / Nyquist -> 1, other w -> 2), accumulates per-thread partials
// in registers across its frequencies, and finally atomically adds them into
// sumsq.  The block that owns f == 0 also writes the DC component (spatial
// sums).  Constraints (checked host-side): B*I <= MAX_BI, I <= MAX_I,
// O % 32 == 0, 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[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;

    for (int f = blockIdx.x; f < F; f += gridDim.x) {
        const int w = f - (f / Wh) * Wh;
        const float mult = ((w == 0) || (w == nyq)) ? inv_n : 2.0f * inv_n;

        const long fbaseX = (long)f * BI;
        for (int t = threadIdx.x; t < BI; t += BLOCK_SIZE)
            sX[t] = X[fbaseX + t];

        const long fbaseW = (long)f * I * O + oc;
        const int IW = I << 5;
        for (int t = threadIdx.x; t < IW; t += BLOCK_SIZE)
            sW[t] = W[fbaseW + (long)(t >> 5) * O + (t & 31)];

        __syncthreads();

        #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 + pb[k] * I;
                const float2* wcol = sW + 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 = sum
                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");
    int64_t F = X.size(0);
    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;
    int f_blocks = (int)std::min<int64_t>(F, (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);
    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);
    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);
    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_v2",
    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

    # ---------------------------------------------------------------- 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

            xp = x.permute(2, 3, 4, 0, 1)                        # [D,H,W,B,In]
            Xf = torch.fft.rfftn(xp, s=s, dim=(0, 1, 2))         # [Do,Ho,Wh,B,In]
            Xf = Xf.reshape(F, B, In)
            Wf = self._get_weight_spectrum(s)                    # [F, In, Out]

            sums, sumsq = _ext.freq_gemm_parseval(
                Xf, 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)
```

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.
