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


_CPP = r"""
#include <torch/extension.h>

torch::Tensor spectral_bn_pool(
    torch::Tensor input_spectrum,
    torch::Tensor weight_spectrum,
    torch::Tensor conv_bias,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches_tracked,
    int64_t spatial,
    int64_t width,
    double scale,
    double eps,
    double momentum,
    bool cumulative,
    bool training);
"""

_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <c10/util/complex.h>
#include <cuda_runtime.h>
#include <cstdint>

constexpr int OUTPUT_TILE = 128;
constexpr int BATCH_TILE = 16;
constexpr int INPUT_TILE = 32;
constexpr int FREQUENCIES_PER_BLOCK = 4;

// Compute Fourier-domain convolutions and immediately accumulate their
// squared magnitudes. No spatial convolution output is materialized.
__global__ void spectral_energy_kernel(
    const float2* __restrict__ inputs,
    const float2* __restrict__ weights,
    float* __restrict__ partial_energy,
    float* __restrict__ sample_means,
    int batch,
    int cin,
    int cout,
    int frequencies,
    int width,
    float inverse_spatial)
{
    __shared__ float2 input_tile[INPUT_TILE][BATCH_TILE];

    const int channel = blockIdx.y * OUTPUT_TILE + threadIdx.x;
    const int batch_start = blockIdx.z * BATCH_TILE;
    const int frequency_start = blockIdx.x * FREQUENCIES_PER_BLOCK;

    float energy[BATCH_TILE];

    #pragma unroll
    for (int n = 0; n < BATCH_TILE; ++n)
        energy[n] = 0.0f;

    for (int df = 0; df < FREQUENCIES_PER_BLOCK; ++df) {
        const int frequency = frequency_start + df;
        if (frequency >= frequencies)
            break;

        float real[BATCH_TILE];
        float imag[BATCH_TILE];

        #pragma unroll
        for (int n = 0; n < BATCH_TILE; ++n) {
            real[n] = 0.0f;
            imag[n] = 0.0f;
        }

        for (int start = 0; start < cin; start += INPUT_TILE) {
            for (int index = threadIdx.x;
                 index < INPUT_TILE * BATCH_TILE;
                 index += OUTPUT_TILE) {
                const int k = index / BATCH_TILE;
                const int n = index % BATCH_TILE;
                const int input_channel = start + k;
                const int sample = batch_start + n;

                float2 value = make_float2(0.0f, 0.0f);
                if (input_channel < cin && sample < batch) {
                    value = inputs[
                        (int64_t(frequency) * cin + input_channel)
                            * batch + sample];
                }
                input_tile[k][n] = value;
            }
            __syncthreads();

            if (channel < cout) {
                const int count = min(INPUT_TILE, cin - start);
                for (int k = 0; k < count; ++k) {
                    const float2 w = weights[
                        (int64_t(frequency) * cin + start + k)
                            * cout + channel];

                    #pragma unroll
                    for (int n = 0; n < BATCH_TILE; ++n) {
                        const float2 a = input_tile[k][n];
                        real[n] = fmaf(a.x, w.x, real[n]);
                        real[n] = fmaf(-a.y, w.y, real[n]);
                        imag[n] = fmaf(a.x, w.y, imag[n]);
                        imag[n] = fmaf(a.y, w.x, imag[n]);
                    }
                }
            }
            __syncthreads();
        }

        if (channel < cout) {
            const int frequency_w = frequency % (width / 2 + 1);
            const float multiplicity =
                (frequency_w == 0 ||
                 ((width & 1) == 0 && frequency_w == width / 2))
                    ? 1.0f : 2.0f;

            #pragma unroll
            for (int n = 0; n < BATCH_TILE; ++n) {
                const float magnitude =
                    fmaf(real[n], real[n], imag[n] * imag[n]);
                energy[n] = fmaf(multiplicity, magnitude, energy[n]);

                if (frequency == 0 && batch_start + n < batch) {
                    sample_means[
                        int64_t(batch_start + n) * cout + channel] =
                            real[n] * inverse_spatial;
                }
            }
        }
    }

    if (channel < cout) {
        #pragma unroll
        for (int n = 0; n < BATCH_TILE; ++n) {
            if (batch_start + n < batch) {
                partial_energy[
                    (int64_t(blockIdx.x) * batch + batch_start + n)
                        * cout + channel] = energy[n];
            }
        }
    }
}

__global__ void reduce_energy_kernel(
    const float* __restrict__ partial_energy,
    double* __restrict__ second_moments,
    int groups,
    int batch,
    int channels,
    double inverse_spatial_squared)
{
    __shared__ double sums[8][32];

    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int channel = blockIdx.x * 32 + lane;
    const int sample = blockIdx.y;

    double sum = 0.0;
    if (channel < channels) {
        for (int group = warp; group < groups; group += 8) {
            sum += double(partial_energy[
                (int64_t(group) * batch + sample) * channels + channel]);
        }
    }

    sums[warp][lane] = sum;
    __syncthreads();

    if (warp == 0 && channel < channels) {
        double total = 0.0;
        #pragma unroll
        for (int w = 0; w < 8; ++w)
            total += sums[w][lane];

        second_moments[int64_t(sample) * channels + channel] =
            total * inverse_spatial_squared;
    }
}

__global__ void finish_bn_pool_kernel(
    const float* __restrict__ sample_means,
    const double* __restrict__ second_moments,
    const float* __restrict__ conv_bias,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int64_t* __restrict__ batches_tracked,
    float* __restrict__ output,
    int batch,
    int channels,
    int64_t spatial,
    float scale,
    float eps,
    float momentum,
    bool cumulative,
    bool training)
{
    __shared__ float update_factor;

    if (threadIdx.x == 0) {
        update_factor = momentum;
        if (training) {
            const int64_t count = ++batches_tracked[0];
            if (cumulative)
                update_factor = 1.0f / float(count);
        }
    }
    __syncthreads();

    for (int channel = threadIdx.x;
         channel < channels;
         channel += blockDim.x) {
        double mean = 0.0;
        double second_moment = 0.0;

        if (training) {
            for (int n = 0; n < batch; ++n) {
                const int64_t index = int64_t(n) * channels + channel;
                mean += double(sample_means[index]);
                second_moment += second_moments[index];
            }
            mean /= double(batch);
            second_moment /= double(batch);

            const double variance =
                fmax(second_moment - mean * mean, 0.0);
            const double scaled_variance =
                double(scale) * double(scale) * variance;
            const float inverse_std =
                rsqrtf(float(scaled_variance) + eps);

            const float scaled_mean =
                float(double(scale) * (mean + double(conv_bias[channel])));
            const double count = double(batch) * double(spatial);
            const float unbiased_variance =
                float(scaled_variance * count / (count - 1.0));

            running_mean[channel] =
                (1.0f - update_factor) * running_mean[channel]
                + update_factor * scaled_mean;
            running_var[channel] =
                (1.0f - update_factor) * running_var[channel]
                + update_factor * unbiased_variance;

            for (int n = 0; n < batch; ++n) {
                const int64_t index = int64_t(n) * channels + channel;
                const float centered =
                    float((double(sample_means[index]) - mean)
                          * double(scale));
                output[index] =
                    centered * inverse_std * gamma[channel] + beta[channel];
            }
        } else {
            const float inverse_std =
                rsqrtf(running_var[channel] + eps);

            for (int n = 0; n < batch; ++n) {
                const int64_t index = int64_t(n) * channels + channel;
                const float value =
                    (sample_means[index] + conv_bias[channel]) * scale;
                output[index] =
                    (value - running_mean[channel])
                    * inverse_std * gamma[channel] + beta[channel];
            }
        }
    }
}

torch::Tensor spectral_bn_pool(
    torch::Tensor input_spectrum,
    torch::Tensor weight_spectrum,
    torch::Tensor conv_bias,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches_tracked,
    int64_t spatial,
    int64_t width,
    double scale,
    double eps,
    double momentum,
    bool cumulative,
    bool training)
{
    const c10::cuda::CUDAGuard guard(input_spectrum.device());

    const int frequencies = input_spectrum.size(0);
    const int cin = input_spectrum.size(1);
    const int batch = input_spectrum.size(2);
    const int channels = weight_spectrum.size(2);
    const int groups =
        (frequencies + FREQUENCIES_PER_BLOCK - 1)
        / FREQUENCIES_PER_BLOCK;

    TORCH_CHECK(
        !training || int64_t(batch) * spatial > 1,
        "Expected more than one value per channel when training");

    auto options = conv_bias.options();
    auto partial = torch::empty({groups, batch, channels}, options);
    auto means = torch::empty({batch, channels}, options);
    auto moments = torch::empty(
        {batch, channels}, options.dtype(torch::kFloat64));
    auto output = torch::empty({batch, channels, 1, 1, 1}, options);

    auto stream = at::cuda::getCurrentCUDAStream();

    spectral_energy_kernel<<<
        dim3(
            groups,
            (channels + OUTPUT_TILE - 1) / OUTPUT_TILE,
            (batch + BATCH_TILE - 1) / BATCH_TILE),
        OUTPUT_TILE, 0, stream>>>(
            reinterpret_cast<const float2*>(
                input_spectrum.data_ptr<c10::complex<float>>()),
            reinterpret_cast<const float2*>(
                weight_spectrum.data_ptr<c10::complex<float>>()),
            partial.data_ptr<float>(),
            means.data_ptr<float>(),
            batch, cin, channels, frequencies, int(width),
            float(1.0 / double(spatial)));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    reduce_energy_kernel<<<
        dim3((channels + 31) / 32, batch), 256, 0, stream>>>(
            partial.data_ptr<float>(),
            moments.data_ptr<double>(),
            groups, batch, channels,
            1.0 / (double(spatial) * double(spatial)));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_bn_pool_kernel<<<1, 256, 0, stream>>>(
        means.data_ptr<float>(),
        moments.data_ptr<double>(),
        conv_bias.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        running_mean.data_ptr<float>(),
        running_var.data_ptr<float>(),
        batches_tracked.data_ptr<int64_t>(),
        output.data_ptr<float>(),
        batch, channels, spatial,
        float(scale), float(eps), float(momentum),
        cumulative, training);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="spectral_transpose3d_bn_pool_v4",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=["spectral_bn_pool"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            verbose=False,
        )
    return _extension


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        scale_factor,
        eps=1e-5,
        momentum=0.1,
    ):
        super().__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))
        self._spectral_cache = None

    def forward(self, x):
        conv = self.conv_transpose
        bn = self.batch_norm
        weight = conv.weight

        tensors = (
            weight,
            conv.bias,
            bn.weight,
            bn.bias,
            bn.running_mean,
            bn.running_var,
        )
        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 5
            and x.numel() > 0
            and x.shape[1] == conv.in_channels
            and not torch.is_autocast_enabled()
            and conv.groups == 1
            and conv.stride == (1, 1, 1)
            and conv.padding == (0, 0, 0)
            and conv.dilation == (1, 1, 1)
            and conv.output_padding == (0, 0, 0)
            and bn.num_batches_tracked is not None
            and bn.num_batches_tracked.device == x.device
            and all(
                tensor is not None
                and tensor.device == x.device
                and tensor.dtype == torch.float32
                and tensor.is_contiguous()
                for tensor in tensors
            )
        )

        if not supported:
            return self.global_avg_pool(bn(conv(x) * self.scale_factor))

        output_shape = tuple(
            size + kernel - 1
            for size, kernel in zip(x.shape[2:], conv.kernel_size)
        )
        spatial = output_shape[0] * output_shape[1] * output_shape[2]

        extension = _get_extension()

        # Only the parameter transform is cached; input statistics are
        # recomputed on every forward pass.
        cache_key = (
            id(weight),
            weight.data_ptr(),
            weight._version,
            tuple(weight.shape),
            tuple(weight.stride()),
            weight.device,
            output_shape,
        )

        with torch.no_grad():
            cache = self._spectral_cache
            if cache is None or cache[0] != cache_key:
                spectrum = torch.fft.rfftn(
                    weight,
                    s=output_shape,
                    dim=(2, 3, 4),
                )
                packed_weight = (
                    spectrum.reshape(conv.in_channels, conv.out_channels, -1)
                    .permute(2, 0, 1)
                    .contiguous()
                )
                self._spectral_cache = (
                    cache_key,
                    weight,
                    packed_weight,
                )
                del spectrum
            else:
                packed_weight = cache[2]

            input_spectrum = torch.fft.rfftn(
                x,
                s=output_shape,
                dim=(2, 3, 4),
            )
            packed_input = (
                input_spectrum.reshape(x.shape[0], conv.in_channels, -1)
                .permute(2, 1, 0)
                .contiguous()
            )
            del input_spectrum

            # DC coefficients give spatial means; Parseval's identity
            # gives second moments without an inverse FFT.
            return extension.spectral_bn_pool(
                packed_input,
                packed_weight,
                conv.bias,
                bn.weight,
                bn.bias,
                bn.running_mean,
                bn.running_var,
                bn.num_batches_tracked,
                spatial,
                output_shape[2],
                float(self.scale_factor),
                float(bn.eps),
                float(bn.momentum) if bn.momentum is not None else 0.0,
                bn.momentum is None,
                bn.training,
            )
```