```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 ConvTranspose3d + BatchNorm3d + AvgPool3d(2) + AvgPool3d(2)
#
# Training path (3 kernels):
#   1. fused_train_kernel: computes the raw ConvTranspose3d output over the
#      full 63^3 volume, fusing (a) 4x4x4 block-sum pooling for the interior
#      60^3 region (written out pre-BN) and (b) per-channel sum / sum-of-
#      squares accumulation for BatchNorm.  Threads are grouped so that each
#      warp handles a single parity class (uniform tap counts, no divergent
#      tap loops) and every thread accumulates R=4 output voxels x 16 output
#      channels to amortize weight loads (weights are staged in shared memory
#      in a float4-friendly [tap][ic][oc] layout).
#   2. bn_finalize_kernel: derives mean/invstd from the accumulated moments,
#      updates running statistics, and emits scale/shift with the 1/64 pool
#      factor folded in.
#   3. bn_apply_kernel: out = raw * scale + shift (vectorized float4).
#
# Inference path (1 kernel):
#   fused_eval_kernel exploits the algebraic commutation of the transposed
#   convolution with the two pooling layers: each pooled output voxel is a
#   3x3x3-tap stencil over 2x2x2 box sums of the input (~8x less arithmetic),
#   with the BatchNorm affine folded into a single scale/shift.  The 63^3
#   intermediate is never formed.
# ============================================================================

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define CI 3
#define CO 16
#define IDIM 32
#define ODIM 15
#define NT 256
#define SXP 5
#define SY 25
#define SXC (CI * SXD)

// Eval geometry: 8x8x4 output tile -> 17x17x9 input slab per image.
#define SLD 17
#define SLH 17
#define SLW 9
#define SLICE (SLD * SLH * SLW)   // 2601
#define PADSL 16
#define SLPAD (SLICE + PADSL)     // 2617

// ---------------------------------------------------------------------------
// Training kernel.
// Block: 256 threads = 8 warps, one warp per parity class (classes ordered by
// tap count).  Tile: 8(d) x 8(h) x 16(w) output voxels -> per class 4x4x8
// voxels -> 32 lanes x R=4 voxels each.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NT)
fused_train_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,      // [27][CI][CO] relaid
    const float* __restrict__ bias,
    float* __restrict__ out,
    double* __restrict__ gsum,
    double* __restrict__ gsq,
    int total_tiles)
{
    __shared__ float s_w[27 * CI * CO];
    __shared__ float s_x[SXC];
    __shared__ float s_red[64][CO];
    __shared__ float s_wsum[8][CO];
    __shared__ float s_wsq[8][CO];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;

    for (int i = tid; i < 27 * CI * CO; i += NT) s_w[i] = w[i];

    // Parity of this warp (class index ordered by tap count).
    const int pd = (warp >= 3 && warp != 6) ? 1 : 0;   // classes 3,4,5,7
    const int ph = (warp == 2 || warp >= 5) ? 1 : 0;   // classes 2,5,6,7
    const int pw = (warp & 1) ? 1 : 0;                 // classes 1,3,5,7

    // Lane-local voxel coordinates within the tile.
    const int ah = lane >> 3;        // 0..3
    const int aw = lane & 7;         // 0..7

    // Tap tables for this class.
    int kd[2], kh[2], kw[2];
    if (pd) { kd[0] = 2; kd[1] = 0; } else { kd[0] = 1; kd[1] = 1; }
    if (ph) { kh[0] = 2; kh[1] = 0; } else { kh[0] = 1; kh[1] = 1; }
    if (pw) { kw[0] = 2; kw[1] = 0; } else { kw[0] = 1; kw[1] = 1; }
    const int nd = pd + 1, nh = ph + 1, nw = pw + 1;
    const int ntaps = nd * nh * nw;

    float tsum[CO], tsq[CO];
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) { tsum[oc] = 0.0f; tsq[oc] = 0.0f; }

    // Voxel flat ids (within the 4x4x8 class subgrid) handled by this lane.
    int vid[4];
    #pragma unroll
    for (int r = 0; r < 4; ++r) vid[r] = (r * 4 + ah) * 8 + aw;

    const int TPI = 4 * 8 * 8;   // tiles per image: (w,h,d)
    __syncthreads();

    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int tw  = rem >> 6;
        rem          &= 63;
        const int th  = rem >> 3;
        const int td  = rem & 7;

        const int a0 = 4 * td, b0 = 4 * th, c0 = 8 * tw;

        // Cooperative load of the 3 x 5 x 5 x 9 input patch (zero padding).
        {
            const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
            for (int i = tid; i < 45; i += NT) {
                const int ic = i / 15;
                const int q  = i - 15 * ic;
                const int a  = q / 3;
                const int bb = q - 3 * a;
                const int ga = a0 + a, gb = b0 + bb;
                const bool ok = (ga < IDIM) && (gb < IDIM);
                float* dst = s_x + ic * SXD + a * 45 + bb * 9;
                if (ok) {
                    const float* src = xb + ((ic * IDIM + ga) * IDIM + gb) * IDIM + c0;
                    #pragma unroll
                    for (int u = 0; u < 9; ++u) {
                        const int gc = c0 + u;
                        dst[u] = (gc < IDIM) ? src[u] : 0.0f;
                    }
                } else {
                    #pragma unroll
                    for (int u = 0; u < 9; ++u) dst[u] = 0.0f;
                }
            }
        }
        __syncthreads();

        float acc[4][CO];
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            const bool val = (8 * td + 2 * r + pd) < 63 &&
                             (8 * th + 2 * ah + ph) < 63 &&
                             (16 * tw + 2 * aw + pw) < 63;
            const float b0f = val ? 1.0f : 0.0f;
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[r][oc] = bias[oc] * b0f;
        }

        for (int tap = 0; tap < ntaps; ++tap) {
            const int od = tap / (nh * nw);
            const int rr = tap - od * (nh * nw);
            const int oh = rr / nw;
            const int ow = rr - oh * nw;
            const int ti = (kd[od] * 9 + kh[oh] * 3 + kw[ow]) * (CI * CO);

            const float w00 = s_w[ti +  0], w01 = s_w[ti +  1], w02 = s_w[ti +  2], w03 = s_w[ti +  3];
            const float w10 = s_w[ti + 16], w11 = s_w[ti + 17], w12 = s_w[ti + 18], w13 = s_w[ti + 19];
            const float w20 = s_w[ti + 32], w21 = s_w[ti + 33], w22 = s_w[ti + 34], w23 = s_w[ti + 35];
            const float w04 = s_w[ti +  4], w05 = s_w[ti +  5], w06 = s_w[ti +  6], w07 = s_w[ti +  7];
            const float w14 = s_w[ti + 20], w15 = s_w[ti + 21], w16 = s_w[ti + 22], w17 = s_w[ti + 23];
            const float w24 = s_w[ti + 36], w25 = s_w[ti + 37], w26 = s_w[ti + 38], w27 = s_w[ti + 39];
            const float w08 = s_w[ti +  8], w09 = s_w[ti +  9], w010 = s_w[ti + 10], w011 = s_w[ti + 11];
            const float w18 = s_w[ti + 24], w19 = s_w[ti + 25], w110 = s_w[ti + 26], w111 = s_w[ti + 27];
            const float w28 = s_w[ti + 40], w29 = s_w[ti + 41], w210 = s_w[ti + 42], w211 = s_w[ti + 43];
            const float w012 = s_w[ti + 12], w013 = s_w[ti + 13], w014 = s_w[ti + 14], w015 = s_w[ti + 15];
            const float w112 = s_w[ti + 28], w113 = s_w[ti + 29], w114 = s_w[ti + 30], w115 = s_w[ti + 31];
            const float w212 = s_w[ti + 44], w213 = s_w[ti + 45], w214 = s_w[ti + 46], w215 = s_w[ti + 47];

            #pragma unroll
            for (int r = 0; r < 4; ++r) {
                const float* xs = s_x + (r + od) * 45 + (ah + oh) * 9 + (aw + ow);
                const float x0 = xs[0];
                const float x1 = xs[SXD];
                const float x2 = xs[2 * SXD];
                acc[r][0]  += x0 * w00  + x1 * w10  + x2 * w20;
                acc[r][1]  += x0 * w01  + x1 * w11  + x2 * w21;
                acc[r][2]  += x0 * w02  + x1 * w12  + x2 * w22;
                acc[r][3]  += x0 * w03  + x1 * w13  + x2 * w23;
                acc[r][4]  += x0 * w04  + x1 * w14  + x2 * w24;
                acc[r][5]  += x0 * w05  + x1 * w15  + x2 * w25;
                acc[r][6]  += x0 * w06  + x1 * w16  + x2 * w26;
                acc[r][7]  += x0 * w07  + x1 * w17  + x2 * w27;
                acc[r][8]  += x0 * w08  + x1 * w18  + x2 * w28;
                acc[r][9]  += x0 * w09  + x1 * w19  + x2 * w29;
                acc[r][10] += x0 * w010 + x1 * w110 + x2 * w210;
                acc[r][11] += x0 * w011 + x1 * w111 + x2 * w211;
                acc[r][12] += x0 * w012 + x1 * w112 + x2 * w212;
                acc[r][13] += x0 * w013 + x1 * w113 + x2 * w213;
                acc[r][14] += x0 * w014 + x1 * w114 + x2 * w214;
                acc[r][15] += x0 * w015 + x1 * w115 + x2 * w215;
            }
        }

        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) {
                const float v = acc[r][oc];
                s_red[vid[r]][oc] = v;
                tsum[oc] += v;
                tsq[oc]  += v * v;
            }
        }
        __syncthreads();

        // Reduce the 8x8x16 tile into 2x2x4 = 16 pool blocks (64 voxels each).
        if (tid < NT / 2) {
            const int bd  = (tid >> 3) & 1;
            const int bh  = (tid >> 2) & 1;
            const int bwl = tid & 3;
            const int oc  = tid >> 4;
            const int base = (bd * 4 + bh) * 8 + bwl * 2;
            float s = 0.0f;
            #pragma unroll
            for (int i1 = 0; i1 < 4; ++i1) {
                #pragma unroll
                for (int i2 = 0; i2 < 4; ++i2) {
                    #pragma unroll
                    for (int i3 = 0; i3 < 4; ++i3) {
                        s += s_red[base + i1 * 32 + i2 * 8 + i3][oc];
                    }
                }
            }
            const int gbd = 2 * td + bd;
            if (gbd < ODIM && th < ODIM && (2 * tw + bwl) < ODIM) {
                const long oidx = (((long)n * CO + oc) * ODIM + gbd) * (ODIM * ODIM)
                                + (long)th * ODIM + (2 * tw + bwl);
                out[oidx] = s;
            }
        }
        __syncthreads();
    }

    // Warp-level reduction of the per-thread moments, then one global
    // double-precision atomicAdd per channel per block.
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) {
        float a = tsum[oc], b = tsq[oc];
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            a += __shfl_down_sync(0xffffffffu, a, o);
            b += __shfl_down_sync(0xffffffffu, b, o);
        }
        if (lane == 0) { s_wsum[warp][oc] = a; s_wsq[warp][oc] = b; }
    }
    __syncthreads();
    if (tid < CO) {
        float a = 0.0f, b = 0.0f;
        #pragma unroll
        for (int wpi = 0; wpi < 8; ++wpi) { a += s_wsum[wpi][tid]; b += s_wsq[wpi][tid]; }
        atomicAdd(gsum + tid, (double)a);
        atomicAdd(gsq + tid, (double)b);
    }
}

// ---------------------------------------------------------------------------
// BatchNorm finalization: computes mean/invstd from the accumulated moments,
// updates running statistics, and produces the fused scale/shift that also
// folds the 1/64 average-pool factor.
// ---------------------------------------------------------------------------
__global__ void bn_finalize_kernel(
    const double* __restrict__ gsum,
    const double* __restrict__ gsq,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    long long* __restrict__ nbt,
    float* __restrict__ scale,
    float* __restrict__ shift,
    double m, double momentum, double eps, int track)
{
    const int c = threadIdx.x;
    if (c >= CO) return;
    const double mean = gsum[c] / m;
    double var = gsq[c] / m - mean * mean;
    if (var < 0.0) var = 0.0;
    const double invstd = rsqrt(var + eps);
    if (track) {
        const double var_unb = var * (m / (m - 1.0));
        running_mean[c] = (float)((double)running_mean[c] * (1.0 - momentum) + momentum * mean);
        running_var[c]  = (float)((double)running_var[c] * (1.0 - momentum) + momentum * var_unb);
        if (c == 0) nbt[0] += 1;
    }
    const double g = (double)gamma[c];
    scale[c] = (float)(g * invstd * (1.0 / 64.0));
    shift[c] = (float)((double)beta[c] - mean * g * invstd);
}

// ---------------------------------------------------------------------------
// Apply the fused BatchNorm + pool scale: out = raw * scale[c] + shift[c].
// ---------------------------------------------------------------------------
__global__ void bn_apply_kernel(
    const float* __restrict__ raw,
    const float* __restrict__ scale,
    const float* __restrict__ shift,
    float* __restrict__ out,
    long total4)
{
    const long i4 = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i4 >= total4) return;
    const long base = i4 << 2;
    const int c = (int)((base >> 2) / (ODIM * ODIM * ODIM)) & (CO - 1);
    const float sc = scale[c], sh = shift[c];
    const float4 v = *reinterpret_cast<const float4*>(raw + base);
    float4 o;
    o.x = v.x * sc + sh;
    o.y = v.y * sc + sh;
    o.z = v.z * sc + sh;
    o.w = v.w * sc + sh;
    *reinterpret_cast<float4*>(out + base) = o;
}

// ---------------------------------------------------------------------------
// Inference kernel: fully fused conv + BN + pool via the commuted stencil.
// ---------------------------------------------------------------------------
__global__ void fused_eval_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    const float* __restrict__ shift,
    float* __restrict__ out)
{
    __shared__ float s_x[CI * SLPAD];
    __shared__ float s_w[CI * 27 * CO];

    const int tid = threadIdx.x;
    const int qd0 = (blockIdx.x & 1) << 3;
    const int qh0 = (blockIdx.x & 2) << 2;
    const int qw0 = (blockIdx.x >> 2) << 2;
    const int n   = blockIdx.y;

    const float* __restrict__ xb = x + (long)n * (CI * IDIM * IDIM * IDIM);
    const float* __restrict__ xr = xb + ((long)2 * qd0 * IDIM + 2 * qh0) * IDIM + 2 * qw0;

    for (int i = tid; i < CI * SLD * SLH; i += NT) {
        const int c  = i / (SLD * SLH);
        const int r  = i - c * (SLD * SLH);
        const int a  = r / SLH;
        const int b  = r - a * SLH;
        const float* src = xr + ((long)c * IDIM + a) * IDIM + b;
        float* dst = s_x + c * SLPAD + r * SLW;
        #pragma unroll
        for (int m = 0; m < SLW; ++m) dst[m] = src[m];
    }
    for (int i = tid; i < CI * 27 * CO; i += NT) s_w[i] = w[i];
    __syncthreads();

    const int qd = qd0 + (tid >> 5);
    const int qh = qh0 + ((tid >> 2) & 7);
    const int qw = qw0 + (tid & 3);
    const float* __restrict__ xs =
        s_x + ((tid >> 5) * SLH + ((tid >> 2) & 7)) * SLW + (tid & 3) * 2;

    float acc[CO];
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;

    const float* __restrict__ wp = s_w;
    #pragma unroll
    for (int ic = 0; ic < CI; ++ic, xs += SLPAD) {
        #pragma unroll
        for (int kkd = 0; kkd < 3; ++kkd) {
            const int os = (kkd == 0) ? SLW : 0;
            float A[3][3];
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                #pragma unroll
                for (int m = 0; m < 3; ++m)
                    A[j][m] = xs[os + j * SLW + m] + xs[os + j * SLW + m + SLW];
            }
            #pragma unroll
            for (int kkh = 0; kkh < 3; ++kkh) {
                const int oh = (kkh == 0) ? 1 : 0;
                float B0 = A[oh][0] + A[oh + 1][0];
                float B1 = A[oh][1] + A[oh + 1][1];
                float B2 = A[oh][2] + A[oh + 1][2];
                const float X0 = B1 + B2;
                const float X1 = B0 + B1;
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) acc[oc] += X0 * wp[oc];
                wp += CO;
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) acc[oc] += X1 * wp[oc];
                wp += CO;
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) acc[oc] += X1 * wp[oc];
                wp += CO;
            }
        }
    }

    const long oidx = (((long)n * CO) * ODIM + qd) * (ODIM * ODIM)
                    + (long)qh * ODIM + qw;
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc)
        out[oidx + (long)oc * (ODIM * ODIM * ODIM)] =
            (acc[oc] + bias[oc]) * scale[oc] + shift[oc];
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                          torch::Tensor gsum, torch::Tensor gsq) {
    const int N = x.size(0);
    auto out = torch::empty({N, CO, ODIM, ODIM, ODIM}, x.options());
    const int total = N * TPI_TOTAL;
    int blocks = total < 2048 ? total : 2048;
    fused_train_kernel<<<blocks, NT>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(),
        gsum.data_ptr<double>(), gsq.data_ptr<double>(), total);
    return out;
}

void bn_finalize(torch::Tensor gsum, torch::Tensor gsq, torch::Tensor gamma,
                 torch::Tensor beta, torch::Tensor running_mean,
                 torch::Tensor running_var, torch::Tensor nbt,
                 torch::Tensor scale, torch::Tensor shift,
                 double m, double momentum, double eps, int64_t track) {
    bn_finalize_kernel<<<1, CO>>>(
        gsum.data_ptr<double>(), gsq.data_ptr<double>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        (long long*)nbt.data_ptr<int64_t>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        m, momentum, eps, (int)track);
}

void bn_apply(torch::Tensor raw, torch::Tensor scale, torch::Tensor shift,
              torch::Tensor out) {
    const long total4 = raw.numel() >> 2;
    const int blocks = (int)((total4 + 255) / 256);
    bn_apply_kernel<<<blocks, 256>>>(
        raw.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>(), total4);
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                         torch::Tensor scale, torch::Tensor shift) {
    const int N = x.size(0);
    auto out = torch::empty({N, CO, ODIM, ODIM, ODIM}, x.options());
    dim3 grid(16, N);
    fused_eval_kernel<<<grid, NT>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>());
    return out;
}
"""

cuda_source = cuda_source.replace("TPI_TOTAL", "256")

cpp_source = (
    "torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor gsum, torch::Tensor gsq);\n"
    "void bn_finalize(torch::Tensor gsum, torch::Tensor gsq, torch::Tensor gamma, "
    "torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, "
    "torch::Tensor nbt, torch::Tensor scale, torch::Tensor shift, "
    "double m, double momentum, double eps, int64_t track);\n"
    "void bn_apply(torch::Tensor raw, torch::Tensor scale, torch::Tensor shift, "
    "torch::Tensor out);\n"
    "torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor scale, torch::Tensor shift);"
)

_ext = load_inline(
    name="fused_convT_bn_pool_v4",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_train", "bn_finalize", "bn_apply", "fused_eval"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: ConvTranspose3d + BatchNorm3d + two AvgPool3d
    layers fused into custom CUDA kernels.  Training uses a single fused
    conv+pool+statistics kernel (parity-class warps, 4 voxels x 16 channels of
    register accumulators per thread), a tiny BatchNorm finalization kernel,
    and a vectorized scale/shift kernel.  Inference uses one fully fused
    kernel exploiting the commutation of the transposed convolution with the
    two pooling layers.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
        self.avg_pool2 = nn.AvgPool3d(kernel_size=2)
        self.ext = _ext

    def _fallback(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.avg_pool1(x)
        x = self.avg_pool2(x)
        return x

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

        supported = (
            x.is_cuda and x.dtype == torch.float32 and
            ct.in_channels == 3 and ct.out_channels == 16 and
            ct.kernel_size == (3, 3, 3) and ct.stride == (2, 2, 2) and
            ct.padding == (1, 1, 1) and ct.output_padding == (0, 0, 0) and
            ct.dilation == (1, 1, 1) and ct.groups == 1 and ct.bias is not None and
            tuple(x.shape[-3:]) == (32, 32, 32) and
            bn.affine and bn.momentum is not None and
            (self.training or bn.track_running_stats)
        )
        if not supported:
            return self._fallback(x)

        x = x.contiguous()
        w = ct.weight
        b = ct.bias.contiguous()
        gamma = bn.weight.contiguous()
        beta = bn.bias.contiguous()
        dev = x.device

        if self.training:
            n = x.shape[0]
            m_count = float(n) * 250047.0  # n * 63^3, exactly representable
            # Relayout weights to [tap(kd,kh,kw)][ic][oc] for float4 access.
            wt = w.permute(2, 3, 4, 1, 0).contiguous().view(-1)

            gsum = torch.zeros(16, dtype=torch.float64, device=dev)
            gsq = torch.zeros(16, dtype=torch.float64, device=dev)
            raw = self.ext.fused_train(x, wt, b, gsum, gsq)

            scale = torch.empty(16, dtype=torch.float32, device=dev)
            shift = torch.empty(16, dtype=torch.float32, device=dev)
            track = bn.track_running_stats
            rm = bn.running_mean if track else gamma
            rv = bn.running_var if track else gamma
            self.ext.bn_finalize(gsum, gsq, gamma, beta, rm, rv,
                                 bn.num_batches_tracked, scale, shift,
                                 m_count, float(bn.momentum), float(bn.eps),
                                 1 if track else 0)

            out = torch.empty_like(raw)
            self.ext.bn_apply(raw, scale, shift, out)
            return out
        else:
            w0 = w.contiguous()
            scale = gamma * torch.rsqrt(bn.running_var + bn.eps) * (1.0 / 64.0)
            shift = beta - bn.running_mean * (gamma * torch.rsqrt(bn.running_var + bn.eps))
            return self.ext.fused_eval(x, w0, b, scale.contiguous(), shift.contiguous())


batch_size = 64
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (out_channels, 1, 1, 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, bias_shape]
```