Looking at this architecture, I need to optimize:
1. **ConvTranspose3d** (16→32 ch, k=3, s=2, p=1) — cuDNN handles this well, I'll keep it
2. **BatchNorm3d** — affine transform that can be algebraically fused
3. **Spatial mean subtraction** — memory-bound passes over a 251MB tensor

Key insights:
- BN + mean-subtract simplifies to `out = scale[c] * (x - spatial_mean[b,c])` where `scale = γ/√(var+ε)` (shift cancels!)
- For k=3, s=2, p=1, the spatial sum of conv output equals a weighted sum of the **input** (33MB read instead of 251MB), via tap-sum regions
- Train mode needs batch stats computed with a custom sum/sumsq kernel

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

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

#define FULL_MASK 0xffffffffu

__inline__ __device__ float warpReduceSum(float v) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(FULL_MASK, v, o);
    return v;
}

// Region sums of input: for k=3,s=2,p=1 convtranspose, output spatial sum
// equals sum over input voxels of x * (sum of valid kernel taps for that voxel's region).
// region = rd*9 + rh*3 + rw, where r in {0: first index, 1: interior, 2: last index}.
__global__ void region_sums_kernel(const float* __restrict__ x, float* __restrict__ out,
                                   int C, int D, int H, int W) {
    int b = blockIdx.x / C;
    int c = blockIdx.x % C;
    __shared__ float sreg[27];
    for (int i = threadIdx.x; i < 27; i += blockDim.x) sreg[i] = 0.f;
    __syncthreads();
    const int S = D * H * W;
    const float* base = x + (long)(b * C + c) * S;
    for (int idx = threadIdx.x; idx < S; idx += blockDim.x) {
        int d = idx / (H * W);
        int r = idx % (H * W);
        int h = r / W;
        int w = r % W;
        int rd = (d == 0) ? 0 : ((d == D - 1) ? 2 : 1);
        int rh = (h == 0) ? 0 : ((h == H - 1) ? 2 : 1);
        int rw = (w == 0) ? 0 : ((w == W - 1) ? 2 : 1);
        atomicAdd(&sreg[(rd * 3 + rh) * 3 + rw], base[idx]);
    }
    __syncthreads();
    float* o = out + ((long)(b * C + c)) * 27;
    for (int i = threadIdx.x; i < 27; i += blockDim.x) o[i] = sreg[i];
}

// Per-(b,c) sum and sum of squares (for train-mode batch statistics).
__global__ void bc_stats_kernel(const float* __restrict__ x, float* __restrict__ sums,
                                float* __restrict__ sumsq, long S) {
    const float* p = x + (long)blockIdx.x * S;
    float s = 0.f, sq = 0.f;
    for (long i = threadIdx.x; i < S; i += blockDim.x) {
        float v = p[i];
        s += v;
        sq += v * v;
    }
    s = warpReduceSum(s);
    sq = warpReduceSum(sq);
    __shared__ float ss[32], ssq[32];
    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    if (lane == 0) { ss[wid] = s; ssq[wid] = sq; }
    __syncthreads();
    if (wid == 0) {
        int nw = (blockDim.x + 31) >> 5;
        s  = (lane < nw) ? ss[lane]  : 0.f;
        sq = (lane < nw) ? ssq[lane] : 0.f;
        s = warpReduceSum(s);
        sq = warpReduceSum(sq);
        if (lane == 0) { sums[blockIdx.x] = s; sumsq[blockIdx.x] = sq; }
    }
}

// Fused final op: out[b,c,i] = scale[c] * (x[b,c,i] - mean[b,c])
// Implements: BatchNorm3d followed by subtraction of the spatial mean.
__global__ void scale_sub_mean_kernel(const float* __restrict__ x,
                                      const float* __restrict__ scale,
                                      const float* __restrict__ mean,
                                      float* __restrict__ out, int C, long S) {
    int c = blockIdx.x % C;
    long off = (long)blockIdx.x * S;
    float sc = scale[c];
    float m = mean[blockIdx.x];
    const float* xp = x + off;
    float* op = out + off;
    for (long i = threadIdx.x; i < S; i += blockDim.x) {
        op[i] = sc * (xp[i] - m);
    }
}

torch::Tensor region_sums_cuda(torch::Tensor x) {
    int B = x.size(0), C = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4);
    auto out = torch::zeros({B, C, 27}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    region_sums_kernel<<<B * C, 128, 0, stream>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), C, D, H, W);
    return out;
}

std::vector<torch::Tensor> bc_stats_cuda(torch::Tensor x) {
    int B = x.size(0), C = x.size(1);
    long S = (long)x.size(2) * x.size(3) * x.size(4);
    auto sums = torch::empty({B, C}, x.options());
    auto sumsq = torch::empty({B, C}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();
    bc_stats_kernel<<<B * C, 256, 0, stream>>>(
        x.data_ptr<float>(), sums.data_ptr<float>(), sumsq.data_ptr<float>(), S);
    return {sums, sumsq};
}

torch::Tensor scale_sub_mean_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor mean) {
    int B = x.size(0), C = x.size(1);
    long S = (long)x.size(2) * x.size(3) * x.size(4);
    auto out = torch::empty_like(x);
    auto stream = at::cuda::getCurrentCUDAStream();
    scale_sub_mean_kernel<<<B * C, 256, 0, stream>>>(
        x.data_ptr<float>(), scale.data_ptr<float>(), mean.data_ptr<float>(),
        out.data_ptr<float>(), C, S);
    return out;
}
"""

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor region_sums_cuda(torch::Tensor x);
std::vector<torch::Tensor> bc_stats_cuda(torch::Tensor x);
torch::Tensor scale_sub_mean_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor mean);
"""

_fused = load_inline(
    name="fused_bn_spatial_meansub",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["region_sums_cuda", "bc_stats_cuda", "scale_sub_mean_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d + BatchNorm3d + spatial-mean subtraction, optimized.

    Math: BN followed by spatial-mean subtraction collapses to
        out[b,c,:] = scale[c] * (conv[b,c,:] - spatial_mean(conv)[b,c]),
        scale[c] = gamma[c] / sqrt(var[c] + eps)
    since the BN shift and channel-mean terms cancel in the subtraction.

    For kernel=3, stride=2, padding=1, the spatial mean of the transposed-conv
    output is a linear functional of the *input*: each input voxel contributes
    its value times the sum of the kernel taps that land in-bounds (a 27-region
    partition). This lets us compute output spatial means by reading only the
    (much smaller) input tensor.
    """

    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 = _fused
        self._fast_mean_ok = (
            kernel_size == 3 and stride == 2 and padding == 1
            and self.conv_transpose.groups == 1
            and self.conv_transpose.dilation == (1, 1, 1)
            and self.conv_transpose.output_padding == (0, 0, 0)
        )
        self._cached_T = None

    def _tap_sum_table(self, device, dtype):
        """T[ci, co, region]: sum of kernel taps valid for each boundary region."""
        if self._cached_T is not None and self._cached_T.device == device:
            return self._cached_T
        # per-dim valid tap masks: region 0 (idx 0) drops tap 0, region 2 (idx n-1)
        # drops tap 2, region 1 (interior) keeps all. Output j = 2*idx - 1 + k.
        M = torch.tensor(
            [[0.0, 1.0, 1.0],
             [1.0, 1.0, 1.0],
             [1.0, 1.0, 0.0]], device=device, dtype=dtype
        )
        W = self.conv_transpose.weight  # [ci, co, kd, kh, kw]
        T = torch.einsum("ioklm,ra,sb,tc->iorst", W, M, M, M)
        T = T.reshape(W.size(0), W.size(1), 27).permute(0, 2, 1).contiguous()  # [ci, 27, co]
        self._cached_T = T
        return T

    def _spatial_means_fast(self, x_in, out_S):
        rs = self.fused.region_sums_cuda(x_in)                      # [B, Ci, 27]
        T = self._tap_sum_table(x_in.device, x_in.dtype)            # [Ci, 27, Co]
        means = torch.einsum("bci,cro->bo", rs.reshape(x_in.size(0), -1),
                             T.reshape(-1, T.size(2)))
        return means / float(out_S)

    def forward(self, x):
        bn = self.batch_norm
        conv = self.conv_transpose(x)
        B, C = conv.size(0), conv.size(1)
        S = conv.size(2) * conv.size(3) * conv.size(4)

        training_stats = self.training or not bn.track_running_stats
        if training_stats:
            sums, sumsq = self.fused.bc_stats_cuda(conv)
            ch_sum = sums.sum(dim=0)
            ch_sumsq = sumsq.sum(dim=0)
            n = float(B) * float(S)
            ch_mean = ch_sum / n
            ch_var = ch_sumsq / n - ch_mean * ch_mean
            scale = bn.weight / torch.sqrt(ch_var + bn.eps)
            mean_bc = sums / float(S)
            with torch.no_grad():
                if bn.track_running_stats:
                    mom = bn.momentum
                    if mom is None:
                        mom = 1.0 / float(bn.num_batches_tracked + 1)
                    bn.running_mean.mul_(1.0 - mom).add_(ch_mean, alpha=mom)
                    unb = n / (n - 1.0) if n > 1.0 else 1.0
                    bn.running_var.mul_(1.0 - mom).add_(ch_var * unb, alpha=mom)
                    bn.num_batches_tracked.add_(1)
        else:
            scale = bn.weight / torch.sqrt(bn.running_var + bn.eps)
            if self._fast_mean_ok and x.is_cuda and x.dtype == torch.float32:
                mean_bc = self._spatial_means_fast(x, S)
            else:
                mean_bc = conv.mean(dim=(2, 3, 4))

        return self.fused.scale_sub_mean_cuda(conv, scale.contiguous(), mean_bc.contiguous())


batch_size = 16
in_channels = 16
out_channels = 32
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1


def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding]
```