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

# Fused kernel: BatchNorm affine + spatial mean subtraction.
#
# Key observation: BN output is a per-channel affine transform z = a_c * y + b_c.
# The final op subtracts the spatial mean of z per (n, c):
#   out = z - mean_spatial(z) = a_c * (y - mean_spatial(y))
# The shift b_c (and hence beta and the BN mean) cancel out exactly, so the
# whole BN + demean chain collapses to a single fused kernel that only needs
# the per-channel scale a_c = gamma_c / sqrt(var_c + eps). Each thread block
# handles one (n, c) slice: a block reduction computes the spatial mean, then
# a second pass writes the normalized, de-meaned output.

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

#define BLOCK_THREADS 512

__global__ void fused_bn_spatial_demean_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ var,
    float* __restrict__ out,
    const int C,
    const long long S,
    const float eps)
{
    const long long slice = blockIdx.x;          // slice = n * C + c
    const int c = (int)(slice % (long long)C);
    const float scale = gamma[c] / sqrtf(var[c] + eps);

    const float* __restrict__ xp = x + slice * S;
    float* __restrict__ op = out + slice * S;

    // Pass 1: block reduction to get the spatial sum.
    float acc = 0.f;
    for (long long i = threadIdx.x; i < S; i += BLOCK_THREADS) {
        acc += xp[i];
    }

    __shared__ float warp_sums[BLOCK_THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        acc += __shfl_down_sync(0xffffffffu, acc, off);

    if (lane == 0) warp_sums[wid] = acc;
    __syncthreads();

    if (wid == 0) {
        float v = (lane < (BLOCK_THREADS / 32)) ? warp_sums[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) warp_sums[0] = v;
    }
    __syncthreads();

    const float mean = warp_sums[0] / (float)S;

    // Pass 2: write scale * (x - mean).
    for (long long i = threadIdx.x; i < S; i += BLOCK_THREADS) {
        op[i] = scale * (xp[i] - mean);
    }
}

torch::Tensor fused_bn_spatial_demean(torch::Tensor x, torch::Tensor gamma, torch::Tensor var, double eps)
{
    auto out = torch::empty_like(x);
    const int C = (int)x.size(1);
    const long long S = (long long)x.size(2) * (long long)x.size(3) * (long long)x.size(4);
    const long long nslices = (long long)x.size(0) * (long long)C;

    dim3 grid((unsigned int)nslices);
    dim3 block(BLOCK_THREADS);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    fused_bn_spatial_demean_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        gamma.data_ptr<float>(),
        var.data_ptr<float>(),
        out.data_ptr<float>(),
        C, S, (float)eps);

    return out;
}
"""

_cpp_src = "torch::Tensor fused_bn_spatial_demean(torch::Tensor x, torch::Tensor gamma, torch::Tensor var, double eps);"

_fused_ext = load_inline(
    name="fused_bn_spatial_demean_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_bn_spatial_demean"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _FusedBNDemean(torch.autograd.Function):
    """out[n,c,i] = (gamma_c / sqrt(var_c + eps)) * (y[n,c,i] - mean_spatial(y[n,c]))"""

    @staticmethod
    def forward(ctx, y, gamma, var, eps):
        out = _fused_ext.fused_bn_spatial_demean(y, gamma, var, eps)
        ctx.save_for_backward(y, gamma, var)
        ctx.eps = eps
        return out

    @staticmethod
    def backward(ctx, grad_out):
        y, gamma, var = ctx.saved_tensors
        g = grad_out.contiguous()
        invstd = torch.rsqrt(var + ctx.eps).view(1, -1, 1, 1, 1)
        scale = gamma.view(1, -1, 1, 1, 1) * invstd
        grad_y = (g - g.mean(dim=(2, 3, 4), keepdim=True)) * scale
        mean_y = y.mean(dim=(2, 3, 4), keepdim=True)
        grad_gamma = (g * (y - mean_y) * invstd).sum(dim=(0, 2, 3, 4))
        return grad_y, grad_gamma, None, None


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cuDNN) + fused BatchNorm/spatial-demean custom 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):
        ct = self.conv_transpose
        y = F.conv_transpose3d(
            x, ct.weight, ct.bias, ct.stride, ct.padding,
            ct.output_padding, ct.groups, ct.dilation,
        )
        bn = self.batch_norm

        fused_ok = (
            y.is_cuda
            and y.dtype == torch.float32
            and bn.affine
            and bn.momentum is not None
            and (bn.training or bn.track_running_stats)
        )
        if not fused_ok:
            z = F.batch_norm(
                y, bn.running_mean, bn.running_var, bn.weight, bn.bias,
                bn.training or not bn.track_running_stats, bn.momentum, bn.eps,
            )
            return z - z.mean(dim=(2, 3, 4), keepdim=True)

        if bn.training:
            if torch.is_grad_enabled() and y.requires_grad:
                # Keep exact autograd semantics of training-mode batch norm.
                z = F.batch_norm(
                    y, bn.running_mean, bn.running_var, bn.weight, bn.bias,
                    True, bn.momentum, bn.eps,
                )
                return z - z.mean(dim=(2, 3, 4), keepdim=True)
            bvar, bmean = torch.var_mean(y, dim=(0, 2, 3, 4), unbiased=False)
            if bn.track_running_stats:
                with torch.no_grad():
                    n = y.numel() // y.size(1)
                    bn.running_mean.lerp_(bmean, bn.momentum)
                    bn.running_var.lerp_(bvar * (float(n) / float(n - 1)), bn.momentum)
            norm_var = bvar
        else:
            norm_var = bn.running_var

        return _FusedBNDemean.apply(
            y.contiguous(), bn.weight, norm_var.contiguous(), float(bn.eps)
        )
```