```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

torch.backends.cudnn.benchmark = True

_fused_cpp_source = """
torch::Tensor fused_bn_meansub(torch::Tensor x, torch::Tensor gamma,
                               torch::Tensor running_mean, torch::Tensor running_var,
                               double eps, double momentum, bool use_batch_stats, bool update_running);
"""

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

#define BLOCK_SIZE 256
#define ITEMS_PER_THREAD 4

// Key math: BN(x) = a[c]*x + b[c], and the final op subtracts the per-(n,c)
// spatial mean, so the additive term b[c] cancels exactly:
//   out[n,c,i] = a[c] * (x[n,c,i] - mean_i(x[n,c,:]))
// where a[c] = gamma[c]/sqrt(var[c]+eps) (batch var in train mode, running
// var in eval mode). Likewise the conv bias cancels, so it is skipped.

__global__ void spatial_reduce_kernel(const float* __restrict__ x,
                                      float* __restrict__ sums,
                                      float* __restrict__ sumsq,
                                      const int S) {
    const int seg = blockIdx.y;  // n*C + c
    const long segStart = (long)seg * (long)S;
    const long segEnd = segStart + (long)S;
    long idx = segStart + (long)blockIdx.x * (BLOCK_SIZE * ITEMS_PER_THREAD) + threadIdx.x;

    float acc = 0.0f;
    float acc2 = 0.0f;
#pragma unroll
    for (int k = 0; k < ITEMS_PER_THREAD; ++k) {
        if (idx < segEnd) {
            float v = x[idx];
            acc += v;
            acc2 += v * v;
        }
        idx += BLOCK_SIZE;
    }

    __shared__ float sh_sum[BLOCK_SIZE];
    __shared__ float sh_sumsq[BLOCK_SIZE];
    sh_sum[threadIdx.x] = acc;
    sh_sumsq[threadIdx.x] = acc2;
    __syncthreads();

    for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
        if (threadIdx.x < offset) {
            sh_sum[threadIdx.x] += sh_sum[threadIdx.x + offset];
            sh_sumsq[threadIdx.x] += sh_sumsq[threadIdx.x + offset];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        atomicAdd(&sums[seg], sh_sum[0]);
        atomicAdd(&sumsq[seg], sh_sumsq[0]);
    }
}

__global__ void channel_params_kernel(const float* __restrict__ sums,
                                      const float* __restrict__ sumsq,
                                      const float* __restrict__ gamma,
                                      float* __restrict__ running_mean,
                                      float* __restrict__ running_var,
                                      float* __restrict__ scale,
                                      const int N, const int C, const double S,
                                      const float eps, const float momentum,
                                      const bool use_batch_stats, const bool update_running) {
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;

    float a;
    if (use_batch_stats) {
        const double M = (double)N * S;
        double s = 0.0;
        double sq = 0.0;
        for (int n = 0; n < N; ++n) {
            s += (double)sums[n * C + c];
            sq += (double)sumsq[n * C + c];
        }
        const double mu = s / M;
        double var = sq / M - mu * mu;
        if (var < 0.0) var = 0.0;
        a = gamma[c] / (float)sqrt(var + (double)eps);
        if (update_running) {
            const double m = (double)momentum;
            const double unb = var * (M / (M - 1.0));  // unbiased estimator
            running_mean[c] = (float)((1.0 - m) * (double)running_mean[c] + m * mu);
            running_var[c] = (float)((1.0 - m) * (double)running_var[c] + m * unb);
        }
    } else {
        a = gamma[c] / sqrtf(running_var[c] + eps);
    }
    scale[c] = a;
}

__global__ void apply_scale_meansub_kernel(const float* __restrict__ x,
                                           const float* __restrict__ sums,
                                           const float* __restrict__ scale,
                                           float* __restrict__ out,
                                           const int C, const int S, const float invS) {
    const int seg = blockIdx.y;
    const int c = seg % C;
    const float a = scale[c];
    const float mean = sums[seg] * invS;
    const long segStart = (long)seg * (long)S;
    const long segEnd = segStart + (long)S;
    long idx = segStart + (long)blockIdx.x * (BLOCK_SIZE * ITEMS_PER_THREAD) + threadIdx.x;

#pragma unroll
    for (int k = 0; k < ITEMS_PER_THREAD; ++k) {
        if (idx < segEnd) {
            out[idx] = a * (x[idx] - mean);
        }
        idx += BLOCK_SIZE;
    }
}

torch::Tensor fused_bn_meansub(torch::Tensor x, torch::Tensor gamma,
                               torch::Tensor running_mean, torch::Tensor running_var,
                               double eps, double momentum, bool use_batch_stats, bool update_running) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    long S = 1;
    for (int i = 2; i < (int)x.dim(); ++i) S *= x.size(i);

    auto out = torch::empty_like(x);
    auto fopts = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
    auto sums = torch::zeros({(long)N * (long)C}, fopts);
    auto sumsq = torch::zeros({(long)N * (long)C}, fopts);
    auto scale = torch::empty({(long)C}, fopts);

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

    const int elems_per_block = BLOCK_SIZE * ITEMS_PER_THREAD;
    const int chunks = (int)((S + elems_per_block - 1) / elems_per_block);
    dim3 grid(chunks, N * C);

    spatial_reduce_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        x.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(), (int)S);

    const int cthreads = 128;
    const int cblocks = (C + cthreads - 1) / cthreads;
    channel_params_kernel<<<cblocks, cthreads, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(), gamma.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        scale.data_ptr<float>(), N, C, (double)S, (float)eps, (float)momentum,
        use_batch_stats, update_running);

    const float invS = (float)(1.0 / (double)S);
    apply_scale_meansub_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        x.data_ptr<float>(), sums.data_ptr<float>(), scale.data_ptr<float>(),
        out.data_ptr<float>(), C, (int)S, invS);

    return out;
}
"""

_ext = load_inline(
    name="fused_bn_spatial_meansub",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_bn_meansub"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cuDNN) + fused BatchNorm3d/spatial-mean-subtraction CUDA kernel.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding, bias=bias)
        self.batch_norm = nn.BatchNorm3d(out_channels)

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

        if not (x.is_cuda and x.dtype == torch.float32):
            y = F.conv_transpose3d(x, ct.weight, ct.bias, ct.stride, ct.padding,
                                   ct.output_padding, ct.groups, ct.dilation)
            y = bn(y)
            return y - torch.mean(y, dim=(2, 3, 4), keepdim=True)

        # The conv bias cancels in the final spatial-mean subtraction, so skip it.
        y = F.conv_transpose3d(x, ct.weight, None, ct.stride, ct.padding,
                               ct.output_padding, ct.groups, ct.dilation)

        gamma = bn.weight
        if gamma is None:
            gamma = torch.ones(bn.num_features, device=y.device, dtype=y.dtype)
        running_mean = bn.running_mean
        running_var = bn.running_var
        if running_mean is None:
            running_mean = gamma  # dummy, unused
            running_var = gamma

        use_batch_stats = bn.training or (not bn.track_running_stats)
        update_running = bn.training and bn.track_running_stats
        momentum = bn.momentum
        if momentum is None:
            if update_running:
                bn.num_batches_tracked += 1
                momentum = 1.0 / float(bn.num_batches_tracked.item())
            else:
                momentum = 0.1

        return _ext.fused_bn_meansub(y, gamma, running_mean, running_var,
                                     float(bn.eps), float(momentum),
                                     use_batch_stats, update_running)
```