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

# Turn-6 refinements (tail kernels were already ~at the HBM roofline; conv is
# left to cuDNN's indexed strided-dgrad which is compute-bound at ~47% of TF32
# peak -- no realistic gain from a hand-rolled replacement):
#  1) Custom float4 NCHW->NHWC conversion kernel for the conv input (replaces
#     the slower at::native permute-copy).
#  2) The channels_last weight copy is cached and only rebuilt when the
#     weight version changes (removes one device kernel + alloc per forward).
#  3) stats kernel: 64-wide spatial tiles (2 float4 loads/thread/iter), more
#     blocks, and a DESCENDING spatial sweep so it starts on the portion of the
#     conv output most likely still resident in L2.
#  4) apply kernel keeps its ASCENDING sweep so its first reads hit the lines
#     most recently touched by the stats pass (L2 hand-off), and its global
#     loads switch to __ldg (evict-normal) so those hits are retained.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256

// ---------------------------------------------------------------------------
// NCHW fallback kernels (layout-generic path)
// ---------------------------------------------------------------------------
__global__ void bn_ms_stats_kernel(
    const float* __restrict__ y,
    double* __restrict__ slice_sum,
    double* __restrict__ chan_sum,
    double* __restrict__ chan_sumsq,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const long base = (long)slice * slice_elems;

    int head = (int)(base & 3L);
    head = head ? (4 - head) : 0;
    if (head > slice_elems) head = slice_elems;
    const int rem  = slice_elems - head;
    const int vec  = rem >> 2;
    const int tail = rem & 3;

    float s = 0.f, sq = 0.f;
    const float4* y4 = reinterpret_cast<const float4*>(y + base + head);
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < vec; i += gridDim.x * THREADS) {
        const float4 v = y4[i];
        s  += v.x + v.y + v.z + v.w;
        sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const float v = y[base + threadIdx.x];
            s += v; sq += v * v;
        }
        if (threadIdx.x < tail) {
            const float v = y[base + head + ((long)vec << 2) + threadIdx.x];
            s += v; sq += v * v;
        }
    }

    __shared__ double bsum[THREADS];
    __shared__ double bsq[THREADS];
    bsum[threadIdx.x] = (double)s;
    bsq[threadIdx.x]  = (double)sq;
    __syncthreads();
    #pragma unroll
    for (int st = THREADS / 2; st > 0; st >>= 1) {
        if (threadIdx.x < st) {
            bsum[threadIdx.x] += bsum[threadIdx.x + st];
            bsq[threadIdx.x]  += bsq[threadIdx.x + st];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        atomicAdd(&slice_sum[slice], bsum[0]);
        atomicAdd(&chan_sum[oc],     bsum[0]);
        atomicAdd(&chan_sumsq[oc],   bsq[0]);
    }
}

__global__ void bn_ms_apply_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,
    const float* __restrict__ slice_mean,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const float sc = scale[oc];
    const float m  = slice_mean[slice];
    const long base = (long)slice * slice_elems;

    int head = (int)(base & 3L);
    head = head ? (4 - head) : 0;
    if (head > slice_elems) head = slice_elems;
    const int rem  = slice_elems - head;
    const int vec  = rem >> 2;
    const int tail = rem & 3;

    const float4* y4 = reinterpret_cast<const float4*>(y + base + head);
    float4* o4 = reinterpret_cast<float4*>(out + base + head);
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < vec; i += gridDim.x * THREADS) {
        const float4 v = y4[i];
        float4 r;
        r.x = sc * (v.x - m);
        r.y = sc * (v.y - m);
        r.z = sc * (v.z - m);
        r.w = sc * (v.w - m);
        o4[i] = r;
    }
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const long idx = base + threadIdx.x;
            out[idx] = sc * (y[idx] - m);
        }
        if (threadIdx.x < tail) {
            const long idx = base + head + ((long)vec << 2) + threadIdx.x;
            out[idx] = sc * (y[idx] - m);
        }
    }
}

// ---------------------------------------------------------------------------
// Shared finalize kernel: per-slice means + per-channel scales (+running stats)
// ---------------------------------------------------------------------------
__global__ void bn_ms_finalize_kernel(
    const double* __restrict__ slice_sum,
    const double* __restrict__ chan_sum,
    const double* __restrict__ chan_sumsq,
    const float* __restrict__ gamma,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ slice_mean,
    float* __restrict__ scale,
    int NSLICE, int OC,
    double inv_spatial, double inv_total, double total,
    float eps, float momentum,
    int use_batch, int update_running)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < NSLICE) {
        slice_mean[i] = (float)(slice_sum[i] * inv_spatial);
    }
    if (i < OC) {
        float sc;
        if (use_batch) {
            const double mu = chan_sum[i] * inv_total;
            double var = chan_sumsq[i] * inv_total - mu * mu;
            if (var < 0.0) var = 0.0;
            sc = gamma[i] * (float)rsqrt(var + (double)eps);
            if (update_running) {
                running_mean[i] = (1.f - momentum) * running_mean[i] + momentum * (float)mu;
                const double unb = var * (total / (total - 1.0));
                running_var[i]  = (1.f - momentum) * running_var[i] + momentum * (float)unb;
            }
        } else {
            sc = gamma[i] * rsqrtf(running_var[i] + eps);
        }
        scale[i] = sc;
    }
}

// ---------------------------------------------------------------------------
// NCHW [N,16,S] -> NHWC [N,S,16] float4 conversion (requires S % 64 == 0).
// Loads: 256B contiguous per half-warp. Stores: 512B contiguous per warp.
// ---------------------------------------------------------------------------
__global__ void nchw_to_nhwc_c16_kernel(
    const float* __restrict__ x,
    float* __restrict__ y, long S)
{
    __shared__ float4 tile[16][17];   // 68-float stride keeps rows 16B aligned
    const int n = blockIdx.y;
    const long s0 = (long)blockIdx.x << 6;
    const int c  = threadIdx.x >> 4;   // 0..15
    const int s4 = threadIdx.x & 15;   // 0..15

    const float4 v = *reinterpret_cast<const float4*>(
        x + ((long)n * 16 + c) * S + s0 + (s4 << 2));
    tile[c][s4] = v;
    __syncthreads();

    const int s  = threadIdx.x >> 2;   // 0..63
    const int cc = (threadIdx.x & 3) << 2;
    const float* t = reinterpret_cast<const float*>(&tile[0][0]);
    float4 o;
    o.x = t[cc * 68 + s];
    o.y = t[(cc + 1) * 68 + s];
    o.z = t[(cc + 2) * 68 + s];
    o.w = t[(cc + 3) * 68 + s];
    *reinterpret_cast<float4*>(y + ((long)n * S + s0 + s) * 16 + cc) = o;
}

// ---------------------------------------------------------------------------
// NHWC (channels_last_3d) kernels. y is logically [N, C, D, H, W] stored as
// [N, S=D*H*W, C] with C == 32 (guarantees 16B alignment of every float4).
// ---------------------------------------------------------------------------

// Pass 1 (NHWC): per-(n,c) spatial sums + per-channel sum / sum-of-squares.
// 64(s) x 32(c) tiles, swept in DESCENDING tile order so the pass starts on
// the most recently produced (possibly L2-resident) part of the conv output.
// Block: 256 threads = 32 spatial lanes x 8 float4 channel groups.
__global__ void stats_nhwc_kernel(
    const float* __restrict__ y,
    double* __restrict__ slice_sum,   // [N*C]
    double* __restrict__ chan_sum,    // [C]
    double* __restrict__ chan_sumsq,  // [C]
    int C, long S, int tiles)
{
    const int n = blockIdx.y;
    const long base = (long)n * S * C;
    const int cg = threadIdx.x & 7;        // 0..7
    const int c0 = cg << 2;                // 0,4,...,28
    const int s_lane = threadIdx.x >> 3;   // 0..31

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;

    for (int t = tiles - 1 - blockIdx.x; t >= 0; t -= gridDim.x) {
        const long ts = (long)t << 6;
        const long sa = ts + s_lane;
        const long sb = sa + 32;
        const float* p = y + base + c0;
        if (sa < S) {
            const float4 v = *reinterpret_cast<const float4*>(p + sa * C);
            s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
            q0 = fmaf(v.x, v.x, q0);
            q1 = fmaf(v.y, v.y, q1);
            q2 = fmaf(v.z, v.z, q2);
            q3 = fmaf(v.w, v.w, q3);
        }
        if (sb < S) {
            const float4 v = *reinterpret_cast<const float4*>(p + sb * C);
            s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
            q0 = fmaf(v.x, v.x, q0);
            q1 = fmaf(v.y, v.y, q1);
            q2 = fmaf(v.z, v.z, q2);
            q3 = fmaf(v.w, v.w, q3);
        }
    }

    // Reduce over the 4 spatial lanes that share a channel group.
    #pragma unroll
    for (int off = 16; off >= 8; off >>= 1) {
        s0 += __shfl_down_sync(0xffffffffu, s0, off);
        s1 += __shfl_down_sync(0xffffffffu, s1, off);
        s2 += __shfl_down_sync(0xffffffffu, s2, off);
        s3 += __shfl_down_sync(0xffffffffu, s3, off);
        q0 += __shfl_down_sync(0xffffffffu, q0, off);
        q1 += __shfl_down_sync(0xffffffffu, q1, off);
        q2 += __shfl_down_sync(0xffffffffu, q2, off);
        q3 += __shfl_down_sync(0xffffffffu, q3, off);
    }

    __shared__ float sh_sum[8][32];
    __shared__ float sh_sq[8][32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane < 8) {
        sh_sum[warp][c0 + 0] = s0; sh_sum[warp][c0 + 1] = s1;
        sh_sum[warp][c0 + 2] = s2; sh_sum[warp][c0 + 3] = s3;
        sh_sq[warp][c0 + 0]  = q0; sh_sq[warp][c0 + 1]  = q1;
        sh_sq[warp][c0 + 2]  = q2; sh_sq[warp][c0 + 3]  = q3;
    }
    __syncthreads();

    if (threadIdx.x < 32) {
        const int c = threadIdx.x;
        float fs = 0.f, fq = 0.f;
        #pragma unroll
        for (int w = 0; w < 8; ++w) {
            fs += sh_sum[w][c];
            fq += sh_sq[w][c];
        }
        const double ds = (double)fs;
        atomicAdd(&slice_sum[n * C + c], ds);
        atomicAdd(&chan_sum[c],          ds);
        atomicAdd(&chan_sumsq[c],        (double)fq);
    }
}

// Pass 2 (NHWC -> NCHW transpose with fused affine):
//   out[n, c, s] = scale[c] * (y[n, s, c] - slice_mean[n, c])
// 64(s) x 32(c) shared-memory tile, ASCENDING sweep so early blocks re-read
// the lines most recently touched by the stats pass (L2 hits). Loads are
// coalesced float4 over channels; each warp stores 32 consecutive spatial
// positions of one channel (128B/warp, fully coalesced).
__global__ void apply_transpose_nhwc_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,      // [C]
    const float* __restrict__ slice_mean, // [N*C]
    int C, long S)
{
    __shared__ float tile[32][65];
    const int n = blockIdx.y;
    const long s0 = (long)blockIdx.x << 6;

    const int cg = threadIdx.x & 7;        // float4 channel group 0..7
    const int s_sub = threadIdx.x >> 3;    // 0..31
    const float4 sc4 = *reinterpret_cast<const float4*>(scale + (cg << 2));
    const float4 mn4 = *reinterpret_cast<const float4*>(slice_mean + n * C + (cg << 2));
    const long nbase = (long)n * S;

    #pragma unroll
    for (int half = 0; half < 2; ++half) {
        const int ss = s_sub + (half << 5);   // 0..63
        const long s = s0 + ss;
        if (s < S) {
            const float4 v = __ldg(reinterpret_cast<const float4*>(
                y + (nbase + s) * C + (cg << 2)));
            tile[(cg << 2) + 0][ss] = sc4.x * (v.x - mn4.x);
            tile[(cg << 2) + 1][ss] = sc4.y * (v.y - mn4.y);
            tile[(cg << 2) + 2][ss] = sc4.z * (v.z - mn4.z);
            tile[(cg << 2) + 3][ss] = sc4.w * (v.w - mn4.w);
        }
    }
    __syncthreads();

    const long rem = S - s0;
    const int s = threadIdx.x & 63;
    if (s < rem) {
        int c = threadIdx.x >> 6;   // 0..3, strides of 4 over 32 channels
        float* obase = out + ((long)n * C) * S + s0 + s;
        #pragma unroll
        for (int i = 0; i < 8; ++i, c += 4) {
            __stcs(obase + (long)c * S, tile[c][s]);
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
static inline void launch_finalize(
    double* slice_sum_p, double* chan_sum_p, double* chan_sumsq_p,
    float* gamma_p, float* rm_p, float* rv_p,
    float* slice_mean_p, float* scale_p,
    int NSLICE, int OC, int slice_elems, int N,
    float eps, float momentum, bool use_batch, bool update_running,
    cudaStream_t stream)
{
    const double spatial = (double)slice_elems;
    const double total   = spatial * (double)N;
    const int work = NSLICE > OC ? NSLICE : OC;
    bn_ms_finalize_kernel<<<(work + 255) / 256, 256, 0, stream>>>(
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        gamma_p, rm_p, rv_p,
        slice_mean_p, scale_p,
        NSLICE, OC, 1.0 / spatial, 1.0 / total, total,
        eps, momentum,
        use_batch ? 1 : 0, update_running ? 1 : 0);
}

// x: [N,16,S] NCHW contiguous -> [N,S,16] NHWC contiguous. Requires S%64==0.
torch::Tensor nchw_to_nhwc_c16(torch::Tensor x)
{
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32 && x.dim() == 5,
                "nchw_to_nhwc_c16 expects 5D CUDA float32 input");
    TORCH_CHECK(x.is_contiguous(), "nchw_to_nhwc_c16 expects NCHW-contiguous input");
    const int N = x.size(0);
    TORCH_CHECK(x.size(1) == 16, "nchw_to_nhwc_c16 requires 16 input channels");
    const long S = x.numel() / ((long)N * 16L);
    TORCH_CHECK((S & 63L) == 0, "nchw_to_nhwc_c16 requires spatial size % 64 == 0");

    auto y = torch::empty({N, 16, x.size(2), x.size(3), x.size(4)}, x.options())
                 .contiguous(at::MemoryFormat::ChannelsLast3d);
    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid((unsigned)(S >> 6), N);
    nchw_to_nhwc_c16_kernel<<<grid, THREADS, 0, stream>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), S);
    return y;
}

// NCHW path (fallback)
torch::Tensor fused_bn_meansub(
    torch::Tensor y,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine)
{
    TORCH_CHECK(y.is_cuda() && y.dtype() == torch::kFloat32,
                "fused_bn_meansub expects CUDA float32 input");
    auto yc = y.contiguous();
    const int N  = yc.size(0);
    const int OC = yc.size(1);
    const long slice_elems_l = yc.numel() / ((long)N * (long)OC);
    const int slice_elems = (int)slice_elems_l;
    const int NSLICE = N * OC;

    auto out = torch::empty_like(yc);
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({OC}, yc.options());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(yc.device());
    auto dbuf = torch::empty({(long)NSLICE + 2L * OC}, opts_d);
    auto fbuf = torch::empty({(long)NSLICE + OC}, yc.options());
    double* slice_sum_p  = dbuf.data_ptr<double>();
    double* chan_sum_p   = slice_sum_p + NSLICE;
    double* chan_sumsq_p = chan_sum_p + OC;
    float*  slice_mean_p = fbuf.data_ptr<float>();
    float*  scale_p      = slice_mean_p + NSLICE;

    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(dbuf.data_ptr(), 0, dbuf.numel() * sizeof(double), stream.stream());

    const int max_vec = (slice_elems + 3) / 4;
    int chunks = (max_vec + THREADS * 4 - 1) / (THREADS * 4);
    if (chunks < 1) chunks = 1;
    if (chunks > 65535) chunks = 65535;
    dim3 grid(chunks, NSLICE);

    bn_ms_stats_kernel<<<grid, THREADS, 0, stream>>>(
        yc.data_ptr<float>(),
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        OC, slice_elems);

    launch_finalize(slice_sum_p, chan_sum_p, chan_sumsq_p,
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean_p, scale_p,
        NSLICE, OC, slice_elems, N,
        (float)eps, (float)momentum, use_batch, update_running,
        stream.stream());

    bn_ms_apply_kernel<<<grid, THREADS, 0, stream>>>(
        yc.data_ptr<float>(), out.data_ptr<float>(),
        scale_p, slice_mean_p,
        OC, slice_elems);

    return out;
}

// NHWC (channels_last_3d) fast path; requires C == 32.
torch::Tensor fused_bn_meansub_cl(
    torch::Tensor y,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine)
{
    TORCH_CHECK(y.is_cuda() && y.dtype() == torch::kFloat32,
                "fused_bn_meansub_cl expects CUDA float32 input");
    TORCH_CHECK(y.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "fused_bn_meansub_cl expects channels_last_3d input");
    const int N  = y.size(0);
    const int OC = y.size(1);
    TORCH_CHECK(OC == 32, "fused_bn_meansub_cl requires 32 channels");
    const long S = y.numel() / ((long)N * (long)OC);
    const int NSLICE = N * OC;

    // final output: NCHW contiguous
    auto out = torch::empty({N, OC, y.size(2), y.size(3), y.size(4)}, y.options());
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({OC}, y.options());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(y.device());
    auto dbuf = torch::empty({(long)NSLICE + 2L * OC}, opts_d);
    auto fbuf = torch::empty({(long)NSLICE + OC}, y.options());
    double* slice_sum_p  = dbuf.data_ptr<double>();
    double* chan_sum_p   = slice_sum_p + NSLICE;
    double* chan_sumsq_p = chan_sum_p + OC;
    float*  slice_mean_p = fbuf.data_ptr<float>();
    float*  scale_p      = slice_mean_p + NSLICE;

    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(dbuf.data_ptr(), 0, dbuf.numel() * sizeof(double), stream.stream());

    const int tiles = (int)((S + 63) >> 6);
    int gx = tiles < 64 ? tiles : 64;
    dim3 sgrid(gx, N);
    stats_nhwc_kernel<<<sgrid, THREADS, 0, stream>>>(
        y.data_ptr<float>(),
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        OC, S, tiles);

    launch_finalize(slice_sum_p, chan_sum_p, chan_sumsq_p,
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean_p, scale_p,
        NSLICE, OC, (int)S, N,
        (float)eps, (float)momentum, use_batch, update_running,
        stream.stream());

    dim3 agrid((unsigned)((S + 63) >> 6), N);
    apply_transpose_nhwc_kernel<<<agrid, THREADS, 0, stream>>>(
        y.data_ptr<float>(), out.data_ptr<float>(),
        scale_p, slice_mean_p,
        OC, S);

    return out;
}
"""

cpp_source = """
torch::Tensor nchw_to_nhwc_c16(torch::Tensor x);
torch::Tensor fused_bn_meansub(
    torch::Tensor y, torch::Tensor gamma,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine);
torch::Tensor fused_bn_meansub_cl(
    torch::Tensor y, torch::Tensor gamma,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine);
"""

fused_op = load_inline(
    name="fused_bn_meansub_v5",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["nchw_to_nhwc_c16", "fused_bn_meansub", "fused_bn_meansub_cl"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cudnn, channels_last_3d, no bias kernel needed) followed
    by a fused custom-CUDA BatchNorm3d + spatial-mean-subtraction tail.

    Exact algebra used:
      BN(y)[n,c,s] - mean_s BN(y) == scale[c] * (g[n,c,s] - mean_s g),
    where y = g + bias[c], scale[c] = gamma[c] / sqrt(var_c + eps). The
    per-channel conv bias, BN beta and BN channel mean all cancel exactly,
    so the conv is run with bias=None and the NHWC conv output is consumed
    directly by the fused kernels (stats + transposing apply).
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        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._w_cl = None
        self._w_cl_key = None

    def _weight_cl(self, w):
        key = (w.data_ptr(), w._version, w.device.index)
        if self._w_cl is None or self._w_cl_key != key:
            self._w_cl = w.detach().contiguous(
                memory_format=torch.channels_last_3d)
            self._w_cl_key = key
        return self._w_cl

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

        if (x.is_cuda and x.dtype == torch.float32
                and bn.momentum is not None and ct.groups == 1):
            track = bn.track_running_stats
            use_batch = bool(bn.training) or (not track)
            update_running = bool(bn.training) and bool(track)

            if track:
                rm = bn.running_mean
                rv = bn.running_var
            else:
                rm = torch.zeros(bn.num_features, device=x.device, dtype=torch.float32)
                rv = torch.ones(bn.num_features, device=x.device, dtype=torch.float32)

            if update_running and bn.num_batches_tracked is not None:
                bn.num_batches_tracked += 1

            gamma = bn.weight if bn.affine else torch.empty(0, device=x.device)
            eps = float(bn.eps)
            mom = float(bn.momentum)

            # Bias cancels exactly in BN->mean-sub, so run the conv without it.
            w_cl = self._weight_cl(ct.weight)
            if x.is_contiguous():
                sin = x.numel() // (x.size(0) * x.size(1))
                if x.size(1) == 16 and (sin & 63) == 0:
                    x_cl = fused_op.nchw_to_nhwc_c16(x)
                else:
                    x_cl = x.contiguous(memory_format=torch.channels_last_3d)
            else:
                x_cl = x.contiguous(memory_format=torch.channels_last_3d)
            y = F.conv_transpose3d(
                x_cl, w_cl, None, ct.stride, ct.padding,
                ct.output_padding, ct.groups, ct.dilation)

            if (y.is_contiguous(memory_format=torch.channels_last_3d)
                    and ct.out_channels == 32):
                return fused_op.fused_bn_meansub_cl(
                    y, gamma, rm, rv, eps, mom,
                    use_batch, update_running, bool(bn.affine))

            # layout fallback: NCHW fused kernels
            return fused_op.fused_bn_meansub(
                y.contiguous(), gamma, rm, rv, eps, mom,
                use_batch, update_running, bool(bn.affine))

        # pure-PyTorch fallback (keeps full bias/BN semantics)
        y = self.conv_transpose(x)
        z = self.batch_norm(y)
        return z - torch.mean(z, dim=(2, 3, 4), keepdim=True)
```