```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 fused CUDA kernel.
#
# Key observation: the network computes
#     y = BatchNorm(conv_transpose(x))
#     out = y - mean_spatial(y)
# BatchNorm is a per-channel affine transform:
#     y[n,c,s] = gamma[c] * (x[n,c,s] - mu[c]) / sqrt(var[c] + eps) + beta[c]
# Subtracting the spatial mean cancels mu and beta exactly, leaving:
#     out[n,c,s] = (x[n,c,s] - mean_s(x[n,c,:])) * scale[c]
# with scale[c] = gamma[c] / sqrt(var[c] + eps).
# So BatchNorm + spatial-mean-subtraction fuse into a single kernel that
# (1) reduces the spatial mean of each (n,c) slice, then
# (2) writes out (x - mean) * scale[c].
fused_bn_demean_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void fused_demean_scale_kernel(const float* __restrict__ x,
                                          float* __restrict__ out,
                                          const float* __restrict__ scale,
                                          int C, long S) {
    const int nc = blockIdx.x;              // one block per (batch, channel) slice
    const int c = nc % C;
    const long base = (long)nc * S;
    const float* __restrict__ xp = x + base;
    float* __restrict__ op = out + base;

    // Pass 1: partial spatial sum (coalesced).
    float sum = 0.f;
    for (long i = threadIdx.x; i < S; i += blockDim.x) {
        sum += xp[i];
    }

    // Block-wide reduction.
    __shared__ float sh[33];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }
    if (lane == 0) sh[wid] = sum;
    __syncthreads();
    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        float v = (lane < nwarps) ? sh[lane] : 0.f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, offset);
        }
        if (lane == 0) sh[32] = v / (float)S;
    }
    __syncthreads();
    const float mean = sh[32];
    const float s = scale[c];

    // Pass 2: subtract spatial mean and apply the BN scale (coalesced).
    for (long i = threadIdx.x; i < S; i += blockDim.x) {
        op[i] = (xp[i] - mean) * s;
    }
}

torch::Tensor fused_bn_demean_cuda(torch::Tensor x, torch::Tensor scale) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(scale.scalar_type() == at::kFloat, "scale must be float32");

    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);

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

    const int threads = 512;
    fused_demean_scale_kernel<<<(int)NC, threads>>>(
        xc.data_ptr<float>(),
        out.data_ptr<float>(),
        scale.data_ptr<float>(),
        C, S);
    return out;
}
"""

fused_bn_demean_cpp_source = (
    "torch::Tensor fused_bn_demean_cuda(torch::Tensor x, torch::Tensor scale);"
)

# Compile the fused (BatchNorm-scale + spatial mean subtraction) kernel.
fused_bn_demean = load_inline(
    name="fused_bn_spatial_demean",
    cpp_sources=fused_bn_demean_cpp_source,
    cuda_sources=fused_bn_demean_source,
    functions=["fused_bn_demean_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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

    Mathematically:
        BN(x) - mean_spatial(BN(x)) == (x - mean_spatial(x)) * scale[c]
        scale[c] = gamma[c] / sqrt(var[c] + eps)
    so the entire BN + demean chain collapses to one fused 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)
        self.fused_op = fused_bn_demean

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

        # Fast fused path: CUDA float32 tensors with standard BN configuration.
        if x.is_cuda and x.dtype == torch.float32 and bn.momentum is not None:
            C = x.shape[1]
            if bn.weight is not None:
                gamma = bn.weight
            else:
                gamma = torch.ones(C, device=x.device, dtype=x.dtype)

            # Use batch statistics when training or when running stats are absent.
            use_batch_stats = bn.training or bn.running_mean is None
            if use_batch_stats:
                var, mean = torch.var_mean(x, dim=(0, 2, 3, 4), unbiased=False)
                if bn.training and bn.running_mean is not None:
                    with torch.no_grad():
                        n = x.numel() // C
                        unbiased_var = var * (float(n) / float(n - 1))
                        bn.running_mean.mul_(1.0 - bn.momentum).add_(mean, alpha=bn.momentum)
                        bn.running_var.mul_(1.0 - bn.momentum).add_(unbiased_var, alpha=bn.momentum)
                        if bn.num_batches_tracked is not None:
                            bn.num_batches_tracked.add_(1)
            else:
                var = bn.running_var

            # beta and the BN mean cancel out against the spatial-mean subtraction.
            scale = (gamma * torch.rsqrt(var + bn.eps)).contiguous()
            return self.fused_op.fused_bn_demean_cuda(x, scale)

        # Safe fallback for any non-standard configuration (exact original ops).
        y = bn(x)
        return y - torch.mean(y, dim=(2, 3, 4), keepdim=True)
```