```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 organized as 8x8x16 voxel tiles.  Each warp owns one
#      parity class (uniform tap counts, no divergent tap loops), each thread
#      accumulates R=4 voxels x 16 output channels in registers so weight
#      reads from shared memory (float4-friendly [tap][ic][oc] layout) are
#      amortized across 4 voxels.  The kernel also fuses:
#        - 4x4x4 block-sum pooling for the interior 60^3 region via an
#          in-register butterfly shuffle reduction + small shared stage, and
#        - per-channel sum / sum-of-squares accumulation for BatchNorm.
#   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 2x average-pool layers: each pooled output is a
#   3x3x3 effective-kernel stencil over the input at stride 2 (K[m] built
#   from w sums: [w1+w2, w0+w1+w2, w0] per dimension), giving ~8x less
#   arithmetic than materializing the 63^3 intermediate.  BatchNorm and the
#   conv bias are folded into a single scale/shift.
# ============================================================================

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 SXD 225              // 5*5*9 input patch per channel
#define SXC (CI * SXD)
#define TPI 256              // 8*8*4 tiles per image
#define SPL (SLD * SLH * SLW)

// Eval geometry: 8x8x4 pooled-output tile -> 17x17x9 input slab per image.
#define SLD 17
#define SLH 17
#define SLW 9

// ---------------------------------------------------------------------------
// Training kernel.
// Block: 256 threads = 8 warps, one warp per parity class.  Warp w handles
// parity (pd,ph,pw) = ((w>>2)&1, (w>>1)&1, w&1).  Lane l -> (ah,aw) =
// (l>>3, l&7); each lane computes voxels (vd,vh,vw) = (2r+pd, 2ah+ph,
// 2aw+pw) for r = 0..3 within the 8x8x16 tile.
// ---------------------------------------------------------------------------
__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_pool[8][16][CO];   // [warp][block = bd*8+bh*4+bw][oc]
    __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];

    const int pd = (warp >> 2) & 1;
    const int ph = (warp >> 1) & 1;
    const int pw = warp & 1;

    const int ah = lane >> 3;        // 0..3
    const int aw = lane & 7;         // 0..7

    // Tap tables for this parity 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;

    // BatchNorm statistics accumulators (sum is derived from pool partials).
    float tsum[CO], tsq[CO];
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) { tsum[oc] = 0.0f; tsq[oc] = 0.0f; }
    __syncthreads();

    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t >> 8;
        int rem       = t & 255;
        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;
            #pragma unroll 3
            for (int i = tid; i < SXC; i += NT) {
                const int ic = i / SXD;
                const int q  = i - ic * SXD;
                const int a  = q / 45;
                const int r2 = q - 45 * a;
                const int bb = r2 / 9;
                const int cc = r2 - 9 * bb;
                const int ga = a0 + a, gb = b0 + bb, gc = c0 + cc;
                float v = 0.0f;
                if ((unsigned)ga < IDIM && (unsigned)gb < IDIM && (unsigned)gc < IDIM)
                    v = xb[((ic * IDIM + ga) * IDIM + gb) * IDIM + gc];
                s_x[i] = v;
            }
        }
        __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 vf = val ? 1.0f : 0.0f;
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[r][oc] = bias[oc] * vf;
        }

        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 float* wt = s_w + (kd[od] * 9 + kh[oh] * 3 + kw[ow]) * (CI * CO);

            float x0[4], x1[4], x2[4];
            {
                const float* xs = s_x + od * 45 + (ah + oh) * 9 + (aw + ow);
                #pragma unroll
                for (int r = 0; r < 4; ++r) {
                    x0[r] = xs[r * 45];
                    x1[r] = xs[r * 45 + SXD];
                    x2[r] = xs[r * 45 + 2 * SXD];
                }
            }
            #pragma unroll
            for (int g4 = 0; g4 < 12; ++g4) {
                const float4 wv = *reinterpret_cast<const float4*>(wt + 4 * g4);
                const int oc = 4 * g4;
                #pragma unroll
                for (int r = 0; r < 4; ++r) {
                    const float xv = (g4 < 4) ? x0[r] : (g4 < 8 ? x1[r] : x2[r]);
                    acc[r][oc]     += xv * wv.x;
                    acc[r][oc + 1] += xv * wv.y;
                    acc[r][oc + 2] += xv * wv.z;
                    acc[r][oc + 3] += xv * wv.w;
                }
            }
        }

        // Per-thread pool partials: blocks (bd, bh, bw), bd from voxel pairs.
        float p0[CO], p1[CO];
        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) {
            p0[oc] = acc[0][oc] + acc[1][oc];
            p1[oc] = acc[2][oc] + acc[3][oc];
            tsq[oc] += acc[0][oc] * acc[0][oc] + acc[1][oc] * acc[1][oc]
                     + acc[2][oc] * acc[2][oc] + acc[3][oc] * acc[3][oc];
        }
        // Butterfly over lane bits 0 (aw) and 3 (ah): sums the 2x2 lane group.
        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) {
            p0[oc] += __shfl_xor_sync(0xffffffffu, p0[oc], 1);
            p1[oc] += __shfl_xor_sync(0xffffffffu, p1[oc], 1);
        }
        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) {
            p0[oc] += __shfl_xor_sync(0xffffffffu, p0[oc], 8);
            p1[oc] += __shfl_xor_sync(0xffffffffu, p1[oc], 8);
        }
        // One lane per (ah,aw) pair-group stages into shared; also fold the
        // partials into the BatchNorm sum accumulator.
        if ((lane & 9) == 0) {
            const int bh = lane >> 4;
            const int bw = (lane & 7) >> 1;
            const int blk = bh * 4 + bw;
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) {
                s_pool[warp][blk][oc]     = p0[oc];
                s_pool[warp][blk + 8][oc] = p1[oc];
                tsum[oc] += p0[oc] + p1[oc];
            }
        }
        __syncthreads();

        // Reduce the 8 warp partials of each pool block and write out.
        {
            const int blk = tid >> 4;      // 0..15
            const int oc  = tid & 15;
            const int bd  = blk >> 3;
            const int bh  = (blk >> 2) & 1;
            const int bw  = blk & 3;
            float s = 0.0f;
            #pragma unroll
            for (int wpi = 0; wpi < 8; ++wpi) s += s_pool[wpi][blk][oc];
            const int gbd = 2 * td + bd;
            const int gbh = 2 * th + bh;
            const int gbw = 4 * tw + bw;
            if (gbd < ODIM && gbh < ODIM && gbw < ODIM) {
                const long oidx = (((long)n * CO + oc) * ODIM + gbd) * (ODIM * ODIM)
                                + (long)gbh * ODIM + gbw;
                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 via the commuted 3x3x3 effective kernel.
// Each block computes an 8x8x4 tile of pooled outputs for one image.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NT)
fused_eval_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,      // effective [27][CI][CO] kernel
    const float* __restrict__ scale,
    const float* __restrict__ shift,
    float* __restrict__ out)
{
    __shared__ float s_x[CI * SPL];
    __shared__ float s_w[27 * CI * 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* xb = x + (long)n * (CI * IDIM * IDIM * IDIM);
    const float* xr = xb + ((long)2 * qd0 * IDIM + 2 * qh0) * IDIM + 2 * qw0;

    #pragma unroll 4
    for (int i = tid; i < CI * SPL; i += NT) {
        const int c  = i / SPL;
        const int r  = i - c * SPL;
        const int a  = r / (SLH * SLW);
        const int r2 = r - a * (SLH * SLW);
        const float* src = xr + ((long)c * IDIM + a) * IDIM + (r2 / SLW);
        s_x[i] = src[r2 - (r2 / SLW) * SLW];
    }
    for (int i = tid; i < 27 * CI * 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);
    if (qd >= ODIM || qh >= ODIM || qw >= ODIM) return;

    const float* xs = s_x + ((2 * (tid >> 5)) * (SLH * SLW))
                    + (((tid >> 2) & 7) * 2) * SLW + (tid & 3) * 2;

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

    const float* wp = s_w;
    #pragma unroll
    for (int ic = 0; ic < CI; ++ic, xs += SPL) {
        #pragma unroll
        for (int md = 0; md < 3; ++md) {
            #pragma unroll
            for (int mh = 0; mh < 3; ++mh) {
                const float x0 = xs[0];
                const float x1 = xs[1];
                const float x2 = xs[2];
                #pragma unroll
                for (int g4 = 0; g4 < 4; ++g4) {
                    const float4 wv = *reinterpret_cast<const float4*>(wp + 4 * g4);
                    acc[4 * g4]     += x0 * wv.x;
                    acc[4 * g4 + 1] += x0 * wv.y;
                    acc[4 * g4 + 2] += x0 * wv.z;
                    acc[4 * g4 + 3] += x0 * wv.w;
                }
                wp += CO;
                #pragma unroll
                for (int g4 = 0; g4 < 4; ++g4) {
                    const float4 wv = *reinterpret_cast<const float4*>(wp + 4 * g4);
                    acc[4 * g4]     += x1 * wv.x;
                    acc[4 * g4 + 1] += x1 * wv.y;
                    acc[4 * g4 + 2] += x1 * wv.z;
                    acc[4 * g4 + 3] += x1 * wv.w;
                }
                wp += CO;
                #pragma unroll
                for (int g4 = 0; g4 < 4; ++g4) {
                    const float4 wv = *reinterpret_cast<const float4*>(wp + 4 * g4);
                    acc[4 * g4]     += x2 * wv.x;
                    acc[4 * g4 + 1] += x2 * wv.y;
                    acc[4 * g4 + 2] += x2 * wv.z;
                    acc[4 * g4 + 3] += x2 * wv.w;
                }
                wp += CO;
            }
            xs += SLH * SLW;
        }
    }

    float* op = out + (((long)n * CO) * ODIM + qd) * (ODIM * ODIM)
              + (long)qh * ODIM + qw;
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc)
        op[(long)oc * (ODIM * ODIM * ODIM)] = acc[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;
    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 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>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>());
    return out;
}
"""

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 scale, torch::Tensor shift);"
)

_ext = load_inline(
    name="fused_convT_bn_pool_v5",
    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, float4 weight broadcasts), 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 (a 3x3x3 effective
    kernel at stride 2), never materializing the 63^3 intermediate.
    """

    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, 0, 1).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:
            # Build the commuted 3x3x3 effective kernel:
            # per dimension K[0] = w1 + w2, K[1] = w0 + w1 + w2, K[2] = w0.
            C = torch.tensor([[0.0, 1.0, 1.0],
                              [1.0, 1.0, 1.0],
                              [1.0, 0.0, 0.0]], device=dev, dtype=torch.float32)
            keff = torch.einsum('ak,bj,cl,iojkl->abcio', C, C, C, w).contiguous().view(-1)
            rs = torch.rsqrt(bn.running_var + bn.eps)
            scale = (gamma * rs * (1.0 / 64.0)).contiguous()
            shift = (beta + (b - bn.running_mean) * gamma * rs).contiguous()
            return self.ext.fused_eval(x, keff, scale, shift)


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