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


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

std::vector<torch::Tensor> pack_spectrum(torch::Tensor spectrum, double scale);

torch::Tensor tensorcore_bn_pool(
    torch::Tensor inputs,
    torch::Tensor weights,
    torch::Tensor input_dc,
    torch::Tensor weight_dc,
    torch::Tensor conv_bias,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches_tracked,
    int64_t cin,
    int64_t channels,
    int64_t spatial,
    int64_t width,
    double scale,
    double eps,
    double momentum,
    bool cumulative);
"""

_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.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>
#include <cstdint>
#include <vector>

namespace wmma = nvcuda::wmma;

// Transpose the spectrum and split each component into two FP16 values.
// Keeping the residual avoids the accuracy loss of a single FP16 transform.
__global__ void pack_kernel(
    const float2* __restrict__ source,
    half* __restrict__ packed,
    float* __restrict__ dc,
    int rows,
    int frequencies,
    int channels,
    int height,
    int width_complex,
    int64_t s0,
    int64_t s1,
    int64_t s2,
    int64_t s3,
    int64_t s4,
    float scale)
{
    __shared__ float2 tile[32][33];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int source_row = blockIdx.y * 32 + ty;
    const int source_frequency = blockIdx.x * 32 + tx;

    #pragma unroll
    for (int j = 0; j < 32; j += 8) {
        const int row = source_row + j;
        float2 value = make_float2(0.0f, 0.0f);

        if (row < rows && source_frequency < frequencies) {
            const int w = source_frequency % width_complex;
            const int h = (source_frequency / width_complex) % height;
            const int d = source_frequency / (width_complex * height);
            const int64_t offset =
                int64_t(row / channels) * s0
                + int64_t(row % channels) * s1
                + int64_t(d) * s2
                + int64_t(h) * s3
                + int64_t(w) * s4;
            value = source[offset];
        }
        tile[ty + j][tx] = value;
    }
    __syncthreads();

    const int output_row = blockIdx.y * 32 + tx;
    const int64_t plane = int64_t(rows) * frequencies;

    #pragma unroll
    for (int j = 0; j < 32; j += 8) {
        const int frequency = blockIdx.x * 32 + ty + j;
        if (output_row < rows && frequency < frequencies) {
            const float2 value = tile[tx][ty + j];
            const float re = value.x * scale;
            const float im = value.y * scale;
            const half re_hi = __float2half_rn(re);
            const half im_hi = __float2half_rn(im);
            const half re_lo = __float2half_rn(re - __half2float(re_hi));
            const half im_lo = __float2half_rn(im - __half2float(im_hi));
            const int64_t offset = int64_t(frequency) * rows + output_row;

            packed[offset] = re_hi;
            packed[plane + offset] = im_hi;
            packed[2 * plane + offset] = re_lo;
            packed[3 * plane + offset] = im_lo;

            if (frequency == 0)
                dc[output_row] = value.x;
        }
    }
}

std::vector<torch::Tensor> pack_spectrum(
    torch::Tensor spectrum,
    double scale)
{
    const c10::cuda::CUDAGuard guard(spectrum.device());

    const int rows = spectrum.size(0) * spectrum.size(1);
    const int frequencies =
        spectrum.size(2) * spectrum.size(3) * spectrum.size(4);

    auto packed = torch::empty(
        {4, frequencies, rows}, spectrum.options().dtype(torch::kFloat16));
    auto dc = torch::empty(
        {rows}, spectrum.options().dtype(torch::kFloat32));

    pack_kernel<<<
        dim3((frequencies + 31) / 32, (rows + 31) / 32),
        dim3(32, 8), 0, at::cuda::getCurrentCUDAStream()>>>(
            reinterpret_cast<const float2*>(
                spectrum.data_ptr<c10::complex<float>>()),
            reinterpret_cast<half*>(packed.data_ptr<at::Half>()),
            dc.data_ptr<float>(),
            rows, frequencies, spectrum.size(1),
            spectrum.size(3), spectrum.size(4),
            spectrum.stride(0), spectrum.stride(1),
            spectrum.stride(2), spectrum.stride(3), spectrum.stride(4),
            float(scale));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return {packed, dc};
}

constexpr int BATCH = 16;
constexpr int CHANNEL_TILE = 128;
constexpr int REDUCTION_CHUNK = 1024;

// Each warp computes all 16 samples for 16 output channels.
// The high-low cross terms retain approximately FP32 input precision.
__global__ void tensorcore_energy_kernel(
    const half* __restrict__ inputs,
    const half* __restrict__ weights,
    float* __restrict__ partial,
    int cin,
    int channels,
    int frequencies,
    int width)
{
    __shared__ float real_output[BATCH][CHANNEL_TILE];
    __shared__ float imag_output[BATCH][CHANNEL_TILE];

    const int frequency = blockIdx.x;
    const int warp = threadIdx.x >> 5;
    const int channel_start = blockIdx.y * CHANNEL_TILE;
    const int warp_channel = channel_start + warp * 16;

    using A = wmma::fragment<
        wmma::matrix_a, 16, 16, 16, half, wmma::row_major>;
    using B = wmma::fragment<
        wmma::matrix_b, 16, 16, 16, half, wmma::row_major>;
    using C = wmma::fragment<
        wmma::accumulator, 16, 16, 16, float>;

    C real_acc, imag_acc;
    wmma::fill_fragment(real_acc, 0.0f);
    wmma::fill_fragment(imag_acc, 0.0f);

    if (warp_channel < channels) {
        const int64_t input_plane = int64_t(frequencies) * BATCH * cin;
        const int64_t weight_plane = int64_t(frequencies) * cin * channels;

        const half* input_base = inputs + int64_t(frequency) * BATCH * cin;
        const half* weight_base =
            weights + int64_t(frequency) * cin * channels + warp_channel;

        for (int k = 0; k < cin; k += 16) {
            A ar, ai, ar_low, ai_low;
            B br, bi, br_low, bi_low;

            wmma::load_matrix_sync(ar, input_base + k, cin);
            wmma::load_matrix_sync(ai, input_base + input_plane + k, cin);
            wmma::load_matrix_sync(
                ar_low, input_base + 2 * input_plane + k, cin);
            wmma::load_matrix_sync(
                ai_low, input_base + 3 * input_plane + k, cin);

            const int64_t offset = int64_t(k) * channels;
            wmma::load_matrix_sync(br, weight_base + offset, channels);
            wmma::load_matrix_sync(
                bi, weight_base + weight_plane + offset, channels);
            wmma::load_matrix_sync(
                br_low, weight_base + 2 * weight_plane + offset, channels);
            wmma::load_matrix_sync(
                bi_low, weight_base + 3 * weight_plane + offset, channels);

            wmma::mma_sync(real_acc, ar, br, real_acc);
            wmma::mma_sync(real_acc, ar, br_low, real_acc);
            wmma::mma_sync(real_acc, ar_low, br, real_acc);

            wmma::mma_sync(imag_acc, ar, bi, imag_acc);
            wmma::mma_sync(imag_acc, ar, bi_low, imag_acc);
            wmma::mma_sync(imag_acc, ar_low, bi, imag_acc);
            wmma::mma_sync(imag_acc, ai, br, imag_acc);
            wmma::mma_sync(imag_acc, ai, br_low, imag_acc);
            wmma::mma_sync(imag_acc, ai_low, br, imag_acc);

            #pragma unroll
            for (int i = 0; i < ai.num_elements; ++i) {
                ai.x[i] = __hneg(ai.x[i]);
                ai_low.x[i] = __hneg(ai_low.x[i]);
            }

            wmma::mma_sync(real_acc, ai, bi, real_acc);
            wmma::mma_sync(real_acc, ai, bi_low, real_acc);
            wmma::mma_sync(real_acc, ai_low, bi, real_acc);
        }

        wmma::store_matrix_sync(
            &real_output[0][warp * 16],
            real_acc, CHANNEL_TILE, wmma::mem_row_major);
        wmma::store_matrix_sync(
            &imag_output[0][warp * 16],
            imag_acc, CHANNEL_TILE, wmma::mem_row_major);
    }
    __syncthreads();

    if (threadIdx.x < CHANNEL_TILE) {
        const int local_channel = threadIdx.x;
        const int channel = channel_start + local_channel;

        if (channel < channels) {
            float energy = 0.0f;
            #pragma unroll
            for (int n = 0; n < BATCH; ++n) {
                const float re = real_output[n][local_channel];
                const float im = imag_output[n][local_channel];
                energy = fmaf(re, re, energy);
                energy = fmaf(im, im, energy);
            }

            const int fw = frequency % (width / 2 + 1);
            const float multiplicity =
                (fw == 0 || (!(width & 1) && fw == width / 2))
                    ? 1.0f : 2.0f;

            partial[int64_t(frequency) * channels + channel] =
                energy * multiplicity;
        }
    }
}

__global__ void reduce_energy_kernel(
    const float* __restrict__ partial,
    double* __restrict__ reduced,
    int frequencies,
    int channels)
{
    __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 start = blockIdx.y * REDUCTION_CHUNK;
    const int end = min(start + REDUCTION_CHUNK, frequencies);

    double sum = 0.0;
    if (channel < channels) {
        for (int f = start + warp; f < end; f += 8)
            sum += double(partial[int64_t(f) * 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];
        reduced[int64_t(blockIdx.y) * channels + channel] = total;
    }
}

__global__ void finish_kernel(
    const float* __restrict__ input_dc,
    const float* __restrict__ weight_dc,
    const double* __restrict__ reduced,
    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 cin,
    int channels,
    int groups,
    int64_t spatial,
    float scale,
    float eps,
    float momentum,
    bool cumulative)
{
    __shared__ float update_factor;
    if (threadIdx.x == 0) {
        const int64_t count = ++batches_tracked[0];
        update_factor = cumulative ? 1.0f / float(count) : momentum;
    }
    __syncthreads();

    for (int channel = threadIdx.x;
         channel < channels;
         channel += blockDim.x) {
        double means[BATCH];
        double mean = 0.0;

        #pragma unroll
        for (int n = 0; n < BATCH; ++n) {
            double dc = 0.0;
            for (int k = 0; k < cin; ++k) {
                dc += double(input_dc[n * cin + k])
                    * double(weight_dc[int64_t(k) * channels + channel]);
            }
            means[n] = dc / double(spatial);
            mean += means[n];
        }
        mean /= double(BATCH);

        double energy = 0.0;
        for (int group = 0; group < groups; ++group)
            energy += reduced[int64_t(group) * channels + channel];

        // Packed inputs are scaled by 1/32 and weights by 256.
        const double second_moment =
            energy / (64.0 * double(BATCH)
                      * double(spatial) * double(spatial));
        const double variance = fmax(second_moment - mean * mean, 0.0);
        const double scaled_variance =
            variance * double(scale) * double(scale);
        const float inverse_std = rsqrtf(float(scaled_variance) + eps);

        const float scaled_mean =
            float((mean + double(conv_bias[channel])) * double(scale));
        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;

        #pragma unroll
        for (int n = 0; n < BATCH; ++n) {
            const float centered =
                float((means[n] - mean) * double(scale));
            output[n * channels + channel] =
                centered * inverse_std * gamma[channel] + beta[channel];
        }
    }
}

torch::Tensor tensorcore_bn_pool(
    torch::Tensor inputs,
    torch::Tensor weights,
    torch::Tensor input_dc,
    torch::Tensor weight_dc,
    torch::Tensor conv_bias,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches_tracked,
    int64_t cin,
    int64_t channels,
    int64_t spatial,
    int64_t width,
    double scale,
    double eps,
    double momentum,
    bool cumulative)
{
    const c10::cuda::CUDAGuard guard(inputs.device());

    const int frequencies = inputs.size(1);
    const int groups =
        (frequencies + REDUCTION_CHUNK - 1) / REDUCTION_CHUNK;

    auto options = conv_bias.options();
    auto partial = torch::empty({frequencies, channels}, options);
    auto reduced = torch::empty(
        {groups, channels}, options.dtype(torch::kFloat64));
    auto output = torch::empty({BATCH, channels, 1, 1, 1}, options);

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

    tensorcore_energy_kernel<<<
        dim3(frequencies, (channels + CHANNEL_TILE - 1) / CHANNEL_TILE),
        256, 0, stream>>>(
            reinterpret_cast<const half*>(inputs.data_ptr<at::Half>()),
            reinterpret_cast<const half*>(weights.data_ptr<at::Half>()),
            partial.data_ptr<float>(),
            int(cin), int(channels), frequencies, int(width));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    reduce_energy_kernel<<<
        dim3((channels + 31) / 32, groups), 256, 0, stream>>>(
            partial.data_ptr<float>(),
            reduced.data_ptr<double>(),
            frequencies, int(channels));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_kernel<<<1, 256, 0, stream>>>(
        input_dc.data_ptr<float>(),
        weight_dc.data_ptr<float>(),
        reduced.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>(),
        int(cin), int(channels), groups, spatial,
        float(scale), float(eps), float(momentum), cumulative);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="spectral_transpose3d_tensorcore_bn_pool_v5",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=["pack_spectrum", "tensorcore_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 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]

        # In evaluation mode, all operations following the convolution are
        # affine, so only the convolution's DC coefficient is needed.
        if not bn.training:
            input_sums = x.sum(dim=(2, 3, 4))
            weight_sums = weight.sum(dim=(2, 3, 4))
            means = (input_sums @ weight_sums) / spatial
            means = (means + conv.bias) * self.scale_factor
            result = (
                (means - bn.running_mean)
                * torch.rsqrt(bn.running_var + bn.eps)
                * bn.weight
                + bn.bias
            )
            return result[:, :, None, None, None]

        fast_training = (
            x.shape[0] == 16
            and conv.in_channels % 16 == 0
            and conv.out_channels % 16 == 0
            and bn.num_batches_tracked is not None
            and bn.num_batches_tracked.device == x.device
            and torch.cuda.get_device_capability(x.device)[0] >= 7
        )

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

        extension = _get_extension()
        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, weight_dc = extension.pack_spectrum(
                    spectrum, 256.0
                )
                self._spectral_cache = (
                    cache_key,
                    weight,
                    packed_weight,
                    weight_dc,
                )
                del spectrum
            else:
                packed_weight, weight_dc = cache[2], cache[3]

            spectrum = torch.fft.rfftn(
                x, s=output_shape, dim=(2, 3, 4)
            )
            packed_input, input_dc = extension.pack_spectrum(
                spectrum, 1.0 / 32.0
            )
            del spectrum

            return extension.tensorcore_bn_pool(
                packed_input,
                packed_weight,
                input_dc,
                weight_dc,
                conv.bias,
                bn.weight,
                bn.bias,
                bn.running_mean,
                bn.running_var,
                bn.num_batches_tracked,
                conv.in_channels,
                conv.out_channels,
                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,
            )
```