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

# Custom CUDA kernels that fuse (scale -> BatchNorm3d -> global average pooling)
# into a single read pass over the ConvTranspose3d output.
#
# Math: GAP(BN(s * x)) per (n, c) collapses to
#   out[n, c] = w_c * (s * mean_spatial(x[n, c]) - mu_c) * inv_std_c + b_c
# so we only need the per-(n, c) spatial sums (and sums of squares when batch
# statistics are required), never materializing the normalized tensor.

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

// Block-wide reduction of two float accumulators (sum and sum of squares).
__device__ __forceinline__ void block_reduce2(float& a, float& b) {
    __shared__ float sa[32];
    __shared__ float sb[32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        a += __shfl_down_sync(0xffffffffu, a, o);
        b += __shfl_down_sync(0xffffffffu, b, o);
    }
    if (lane == 0) { sa[warp] = a; sb[warp] = b; }
    __syncthreads();
    const int nw = (blockDim.x + 31) >> 5;
    if (warp == 0) {
        a = (lane < nw) ? sa[lane] : 0.f;
        b = (lane < nw) ? sb[lane] : 0.f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            a += __shfl_down_sync(0xffffffffu, a, o);
            b += __shfl_down_sync(0xffffffffu, b, o);
        }
    }
}

// One block per (n, c); vectorized float4 loads over the spatial volume.
template <bool STATS>
__global__ void gap_sums_vec_kernel(const float4* __restrict__ x4,
                                    float* __restrict__ sum,
                                    float* __restrict__ sumsq,
                                    const int spatial4) {
    const long long nc = blockIdx.x;
    const float4* p = x4 + nc * (long long)spatial4;
    float acc = 0.f, acc2 = 0.f;
    for (int i = threadIdx.x; i < spatial4; i += blockDim.x) {
        const float4 v = p[i];
        acc += (v.x + v.y) + (v.z + v.w);
        if (STATS) acc2 += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
    }
    block_reduce2(acc, acc2);
    if (threadIdx.x == 0) {
        sum[nc] = acc;
        if (STATS) sumsq[nc] = acc2;
    }
}

// Scalar fallback for spatial sizes not divisible by 4.
template <bool STATS>
__global__ void gap_sums_scalar_kernel(const float* __restrict__ x,
                                       float* __restrict__ sum,
                                       float* __restrict__ sumsq,
                                       const long long spatial) {
    const long long nc = blockIdx.x;
    const float* p = x + nc * spatial;
    float acc = 0.f, acc2 = 0.f;
    for (long long i = threadIdx.x; i < spatial; i += blockDim.x) {
        const float v = p[i];
        acc += v;
        if (STATS) acc2 += v * v;
    }
    block_reduce2(acc, acc2);
    if (threadIdx.x == 0) {
        sum[nc] = acc;
        if (STATS) sumsq[nc] = acc2;
    }
}

// One thread per (n, c): fold scale + batchnorm affine into the pooled mean.
__global__ void scale_bn_gap_finalize_kernel(const float* __restrict__ sum,
                                             const float* __restrict__ sumsq,
                                             const float* __restrict__ weight,
                                             const float* __restrict__ bias,
                                             float* __restrict__ running_mean,
                                             float* __restrict__ running_var,
                                             float* __restrict__ out,
                                             const int N, const int C,
                                             const float scale, const float eps,
                                             const float momentum,
                                             const float inv_spatial,
                                             const float inv_cnt,
                                             const int use_batch,
                                             const int update_stats) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * C) return;
    const int n = idx / C;
    const int c = idx - n * C;
    float mu, inv_std;
    if (use_batch) {
        double s1 = 0.0, s2 = 0.0;
        for (int i = 0; i < N; ++i) {
            s1 += (double)sum[i * C + c];
            s2 += (double)sumsq[i * C + c];
        }
        const float m = (float)(s1 * (double)inv_cnt);
        float v = (float)(s2 * (double)inv_cnt) - m * m;
        v = fmaxf(v, 0.f);
        const float var_bn = scale * scale * v;
        mu = scale * m;
        inv_std = rsqrtf(var_bn + eps);
        if (update_stats && n == 0) {
            const float cnt = 1.f / inv_cnt;
            const float unb = var_bn * cnt / (cnt - 1.f);  // unbiased for running stats
            running_mean[c] = (1.f - momentum) * running_mean[c] + momentum * mu;
            running_var[c]  = (1.f - momentum) * running_var[c]  + momentum * unb;
        }
    } else {
        mu = running_mean[c];
        inv_std = rsqrtf(running_var[c] + eps);
    }
    const float mean_nc = sum[idx] * inv_spatial;
    out[idx] = weight[c] * (scale * mean_nc - mu) * inv_std + bias[c];
}

torch::Tensor fused_scale_bn_gap(torch::Tensor x, double scale,
                                 torch::Tensor weight, torch::Tensor bias,
                                 torch::Tensor running_mean, torch::Tensor running_var,
                                 double eps, double momentum, bool training) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    auto xc = x.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    int64_t spatial = 1;
    for (int i = 2; i < xc.dim(); ++i) spatial *= xc.size(i);

    auto opts = xc.options();
    const bool track = running_mean.defined() && running_mean.numel() > 0;
    const bool use_batch = training || !track;
    const bool update_stats = training && track;

    auto sum_t = torch::empty({N, C}, opts);
    torch::Tensor sumsq_t;
    float* sumsq_ptr = nullptr;
    if (use_batch) {
        sumsq_t = torch::empty({N, C}, opts);
        sumsq_ptr = sumsq_t.data_ptr<float>();
    }
    auto out = torch::empty({N, C, 1, 1, 1}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    const dim3 grid((unsigned int)(N * C));

    if ((spatial % 4) == 0) {
        const int sp4 = (int)(spatial >> 2);
        if (use_batch) {
            gap_sums_vec_kernel<true><<<grid, threads, 0, stream>>>(
                reinterpret_cast<const float4*>(xc.data_ptr<float>()),
                sum_t.data_ptr<float>(), sumsq_ptr, sp4);
        } else {
            gap_sums_vec_kernel<false><<<grid, threads, 0, stream>>>(
                reinterpret_cast<const float4*>(xc.data_ptr<float>()),
                sum_t.data_ptr<float>(), sumsq_ptr, sp4);
        }
    } else {
        if (use_batch) {
            gap_sums_scalar_kernel<true><<<grid, threads, 0, stream>>>(
                xc.data_ptr<float>(), sum_t.data_ptr<float>(), sumsq_ptr, spatial);
        } else {
            gap_sums_scalar_kernel<false><<<grid, threads, 0, stream>>>(
                xc.data_ptr<float>(), sum_t.data_ptr<float>(), sumsq_ptr, spatial);
        }
    }

    float* rm_ptr = track ? running_mean.data_ptr<float>() : nullptr;
    float* rv_ptr = track ? running_var.data_ptr<float>() : nullptr;

    const int fthreads = 256;
    const int fblocks = (int)((N * C + fthreads - 1) / fthreads);
    scale_bn_gap_finalize_kernel<<<fblocks, fthreads, 0, stream>>>(
        sum_t.data_ptr<float>(), sumsq_ptr,
        weight.data_ptr<float>(), bias.data_ptr<float>(),
        rm_ptr, rv_ptr, out.data_ptr<float>(),
        (int)N, (int)C, (float)scale, (float)eps, (float)momentum,
        1.0f / (float)spatial, 1.0f / (float)(N * spatial),
        use_batch ? 1 : 0, update_stats ? 1 : 0);

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_scale_bn_gap(torch::Tensor x, double scale, "
    "torch::Tensor weight, torch::Tensor bias, torch::Tensor running_mean, "
    "torch::Tensor running_var, double eps, double momentum, bool training);"
)

fused_scale_bn_gap_ext = load_inline(
    name="fused_scale_bn_gap_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_scale_bn_gap"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuDNN ConvTranspose3d, and replaces
    (scale -> BatchNorm3d -> global average pooling) with a single fused
    custom CUDA reduction kernel.
    """
    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))

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

        # Fast path: fused custom CUDA kernel.
        if x.is_cuda and x.dtype == torch.float32:
            mom = bn.momentum
            if self.training and bn.track_running_stats:
                if mom is None:
                    mom = 1.0 / float(bn.num_batches_tracked.item() + 1)
                bn.num_batches_tracked.add_(1)
            if mom is None:
                mom = 0.0
            if bn.track_running_stats:
                rm, rv = bn.running_mean, bn.running_var
            else:
                rm = torch.empty(0, device=x.device, dtype=torch.float32)
                rv = rm
            return fused_scale_bn_gap_ext.fused_scale_bn_gap(
                x, self.scale_factor, bn.weight, bn.bias, rm, rv,
                float(bn.eps), float(mom), self.training)

        # Fallback: original PyTorch ops.
        x = x * self.scale_factor
        x = bn(x)
        x = self.global_avg_pool(x)
        return x
```