You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # 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


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # 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

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

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

class Model(nn.Module):
    """
    A model that performs a 3D transposed convolution, followed by batch normalization, 
    two average pooling layers.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(Model, 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)

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


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]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 6 of 6 for this task. Your previous solution (turn 5) was:
```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]
```
Result of evaluating it:
The solution failed to compile. Error output:
```
RuntimeError: Error building extension 'fused_convT_bn_pool_v4': [1/3] /opt/kbv/v/lib/python3.12/site-packages/nvidia/cu13/bin/nvcc -MD -MF cuda.cuda.o.d -DTORCH_EXTENSION_NAME=fused_convT_bn_pool_v4 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /opt/kbv/v/lib/python3.12/site-packages/torch/include -isystem /opt/kbv/v/lib/python3.12/site-packages/torch/include/torch/csrc/api/include -isystem /opt/kbv/v/lib/python3.12/site-packages/nvidia/cu13/include -isystem /usr/include/python3.12 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -O3 -std=c++20 -c /srv/jobs/1791595478975709400-3349698-0/build/fused_convT_bn_pool_v4/cuda.cu -o cuda.cuda.o 
FAILED: [code=2] cuda.cuda.o 
/opt/kbv/v/lib/python3.12/site-packages/nvidia/cu13/bin/nvcc -MD -MF cuda.cuda.o.d -DTORCH_EXTENSION_NAME=fused_convT_bn_pool_v4 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /opt/kbv/v/lib/python3.12/site-packages/torch/include -isystem /opt/kbv/v/lib/python3.12/site-packages/torch/include/torch/csrc/api/include -isystem /opt/kbv/v/lib/python3.12/site-packages/nvidia/cu13/include -isystem /usr/include/python3.12 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -O3 -std=c++20 -c /srv/jobs/1791595478975709400-3349698-0/build/fused_convT_bn_pool_v4/cuda.cu -o cuda.cuda.o 
/srv/jobs/1791595478975709400-3349698-0/build/fused_convT_bn_pool_v4/cuda.cu(42): error: identifier "SXD" is undefined
                float s_x[(3 * SXD)];
                               ^

1 error detected in the compilation of "/srv/jobs/1791595478975709400-3349698-0/build/fused_convT_bn_pool_v4/cuda.cu".
[2/3] c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=fused_convT_bn_pool_v4 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /opt/kbv/v/lib/python3.12/site-packages/torch/include -isystem /opt/kbv/v/lib/python3.12/site-packages/torch/include/torch/csrc/api/include -isystem /opt/kbv/v/lib/python3.12/site-packages/nvidia/cu13/include -isystem /usr/include/python3.12 -fPIC -std=c++20 -c /srv/jobs/1791595478975709400-3349698-0/build/fused_convT_bn_pool_v4/main.cpp -o main.o 
ninja: build stopped: subcommand failed.

```

The fastest correct solution so far (22.547x speedup) is:
```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


# ============================================================================
# Custom CUDA kernels
#
# Optimizations vs. the previous version:
#  1. Training mode no longer runs a separate full-volume statistics kernel.
#     A single fused kernel computes the raw ConvTranspose3d output for the
#     whole 63^3 volume (organized as 8x4x4 voxel tiles), writes the 4x4x4
#     block sums (pooled pre-BN values) for the interior 60^3 region, and
#     accumulates per-channel sum / sum-of-squares for BatchNorm on the fly.
#  2. Threads are assigned to voxels sorted by parity class so that warps are
#     (nearly) uniform in the number of kernel taps they execute.  This
#     removes the divergent-loop serialization of the naive mapping and
#     roughly doubles the conv throughput.
#  3. BatchNorm scale/shift are finalized in a tiny dedicated kernel that
#     also updates the running statistics, and BN+pool scaling is applied
#     with a single addcmul, so training mode adds only ~50us of overhead.
#  4. Inference mode keeps everything in one kernel: ConvTranspose3d +
#     BatchNorm(running stats) + AvgPool(4,4,4), never materializing the
#     63^3 intermediate in DRAM.
# ============================================================================

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

#define CI 3
#define CO 16
#define IDIM 32
#define ODIM 15
#define NTHREADS 128
#define PITCH (CO + 1)
#define SXD 45   // 5*3*3 input patch per channel for an 8x4x4 output tile

// Parity classes ordered by tap count so warps are near-uniform:
// 0:eee(1 tap) 1:eeo(2) 2:eoe(2) 3:oee(2) 4:eoo(4) 5:oeo(4) 6:ooe(4) 7:ooo(8)
__constant__ int c_pd[8] = {0, 0, 0, 1, 0, 1, 1, 1};
__constant__ int c_ph[8] = {0, 0, 1, 0, 1, 0, 1, 1};
__constant__ int c_pw[8] = {0, 1, 0, 0, 1, 1, 0, 1};

// For local coordinate l (parity dependent):
//   global D = base + l ; input d = (D + 1 - kd) / 2
//   even l: 1 tap  -> a = l/2   , kd = 1
//   odd  l: 2 taps -> a = l/2   , kd = 2
//                     a = l/2+1 , kd = 0
__device__ __forceinline__ int setup_taps(int l, int* a, int* k) {
    if ((l & 1) == 0) { a[0] = l >> 1; k[0] = 1; return 1; }
    a[0] = l >> 1; k[0] = 2; a[1] = (l >> 1) + 1; k[1] = 0;
    return 2;
}

__device__ __forceinline__ void conv_acc(
    const float* s_x, const float* s_w,
    const int* ad, const int* ah, const int* aw,
    const int* kd, const int* kh, const int* kw,
    int nd, int nh, int nw, float* acc)
{
    for (int i = 0; i < nd; ++i) {
        for (int j = 0; j < nh; ++j) {
            for (int m = 0; m < nw; ++m) {
                const int xs = ad[i] * 9 + ah[j] * 3 + aw[m];
                const float x0 = s_x[xs];
                const float x1 = s_x[xs + SXD];
                const float x2 = s_x[xs + 2 * SXD];
                const int woff = kd[i] * 9 + kh[j] * 3 + kw[m];
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) {
                    const float* wp = s_w + oc * 27 + woff;
                    acc[oc] += x0 * wp[0] + x1 * wp[CO * 27] + x2 * wp[2 * CO * 27];
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Training kernel: raw ConvTranspose3d over the full 63^3 volume, fused
// 4x4x4 block pooling (raw block sums) + BatchNorm statistics accumulation.
// Tiles: 8x4x4 voxels, grid per image = 8*16*16 tiles (covers 64^3 >= 63^3).
// ---------------------------------------------------------------------------
__global__ void fused_train_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    double* __restrict__ gsum,
    double* __restrict__ gsq,
    int total_tiles)
{
    __shared__ float s_x[CI * SXD];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_bsum[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) s_bsum[tid] = 0.0f;

    const int cls = tid >> 4, r = tid & 15;
    const int pd = c_pd[cls], ph = c_ph[cls], pw = c_pw[cls];
    const int hd = r >> 2, hh = (r >> 1) & 1, hw = r & 1;
    const int vd = 2 * hd + pd, vh = 2 * hh + ph, vw = 2 * hw + pw;
    const int vid = vd * 16 + vh * 4 + vw;

    int ad[2], ah[2], aw[2], kd[2], kh[2], kw[2];
    const int nd = setup_taps(vd, ad, kd);
    const int nh = setup_taps(vh, ah, kh);
    const int nw = setup_taps(vw, aw, kw);

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

    const int TPI = 8 * 16 * 16;
    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int td  = rem >> 8;
        rem          &= 255;
        const int th  = rem >> 4;
        const int tw  = rem & 15;
        const bool interior = (td < 7) && (th < 15) && (tw < 15);

        const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        if (interior) {
            #pragma unroll
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                s_x[i] = xb[((ic * IDIM + 4 * td + a) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
            }
        } else {
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                const int gi = 4 * td + a, gj = 2 * th + b, gk = 2 * tw + c;
                float v = 0.0f;
                if (gi < IDIM && gj < IDIM && gk < IDIM)
                    v = xb[((ic * IDIM + gi) * IDIM + gj) * IDIM + gk];
                s_x[i] = v;
            }
        }
        __syncthreads();

        float acc[CO];
        const bool valid = interior ||
            ((8 * td + vd) < 63 && (4 * th + vh) < 63 && (4 * tw + vw) < 63);
        if (valid) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = bias[oc];
            conv_acc(s_x, s_w, ad, ah, aw, kd, kh, kw, nd, nh, nw, acc);
        } else {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;
        }

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

        if (tid < 32) {
            const int oc = tid & 15, half = tid >> 4;
            float s = 0.0f;
            #pragma unroll 8
            for (int i = 0; i < 64; ++i) s += s_red[half * 64 + i][oc];
            atomicAdd(&s_bsum[oc], s);
            const int bd4 = 2 * td + half;
            if (bd4 < ODIM && th < ODIM && tw < ODIM) {
                const long oidx = (((long)n * CO + oc) * ODIM + bd4) * (ODIM * ODIM)
                                + (long)th * ODIM + tw;
                out[oidx] = s;
            }
        }
    }

    __syncthreads();
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) s_red[tid][oc] = tsq[oc];
    __syncthreads();
    if (tid < CO) {
        float q = 0.0f;
        #pragma unroll 8
        for (int i = 0; i < NTHREADS; ++i) q += s_red[i][tid];
        atomicAdd(gsq + tid, (double)q);
        atomicAdd(gsum + tid, (double)s_bsum[tid]);
    }
}

// ---------------------------------------------------------------------------
// Inference kernel: ConvTranspose3d + BatchNorm(scale/shift) + AvgPool(4)
// fully fused.  Only voxels in the 60^3 region that pooling reads are
// computed.  Tiles per image: 8*15*15.
// ---------------------------------------------------------------------------
__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,
    int total_tiles)
{
    __shared__ float s_x[CI * SXD];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_sc[CO];
    __shared__ float s_sh[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) { s_sc[tid] = scale[tid]; s_sh[tid] = shift[tid]; }

    const int cls = tid >> 4, r = tid & 15;
    const int pd = c_pd[cls], ph = c_ph[cls], pw = c_pw[cls];
    const int hd = r >> 2, hh = (r >> 1) & 1, hw = r & 1;
    const int vd = 2 * hd + pd, vh = 2 * hh + ph, vw = 2 * hw + pw;
    const int vid = vd * 16 + vh * 4 + vw;

    int ad[2], ah[2], aw[2], kd[2], kh[2], kw[2];
    const int nd = setup_taps(vd, ad, kd);
    const int nh = setup_taps(vh, ah, kh);
    const int nw = setup_taps(vw, aw, kw);
    __syncthreads();

    const int TPI = 8 * ODIM * ODIM;
    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int td  = rem / (ODIM * ODIM);
        rem          %= ODIM * ODIM;
        const int th  = rem / ODIM;
        const int tw  = rem % ODIM;
        const bool gd = (td == 7);

        const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        if (!gd) {
            #pragma unroll
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                s_x[i] = xb[((ic * IDIM + 4 * td + a) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
            }
        } else {
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                const int gi = 4 * td + a;
                float v = 0.0f;
                if (gi < IDIM)
                    v = xb[((ic * IDIM + gi) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
                s_x[i] = v;
            }
        }
        __syncthreads();

        float acc[CO];
        if (!gd || vd < 4) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = bias[oc];
            conv_acc(s_x, s_w, ad, ah, aw, kd, kh, kw, nd, nh, nw, acc);
        } else {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;
        }

        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) s_red[vid][oc] = acc[oc];
        __syncthreads();

        if (tid < 32) {
            const int oc = tid & 15, half = tid >> 4;
            const int bd4 = 2 * td + half;
            if (bd4 < ODIM) {
                float s = 0.0f;
                #pragma unroll 8
                for (int i = 0; i < 64; ++i) s += s_red[half * 64 + i][oc];
                const long oidx = (((long)n * CO + oc) * ODIM + bd4) * (ODIM * ODIM)
                                + (long)th * ODIM + tw;
                out[oidx] = s * s_sc[oc] + s_sh[oc];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// BatchNorm finalization: computes mean/var 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);
}

// ---------------------------------------------------------------------------
// 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 * 8 * 16 * 16;
    int blocks = total < 4096 ? total : 4096;
    fused_train_kernel<<<blocks, NTHREADS>>>(
        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;
}

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());
    const int total = N * 8 * ODIM * ODIM;
    int blocks = total < 4096 ? total : 4096;
    fused_eval_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>(), 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);
}
"""

cpp_source = (
    "torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor gsum, torch::Tensor gsq);\n"
    "torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor scale, torch::Tensor shift);\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);"
)

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


class ModelNew(nn.Module):
    """
    Optimized version of Model: ConvTranspose3d + BatchNorm3d + two AvgPool3d
    layers fused into custom CUDA kernels.  A single kernel handles the whole
    forward pass in inference mode; training mode uses one fused kernel for
    conv+pool+statistics, a tiny finalization kernel for the BatchNorm
    moments, and one elementwise scale.
    """

    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.contiguous()
        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
            gsum = torch.zeros(16, dtype=torch.float64, device=dev)
            gsq = torch.zeros(16, dtype=torch.float64, device=dev)
            raw = self.ext.fused_train(x, w, 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
            nbt = bn.num_batches_tracked if track else bn.num_batches_tracked
            self.ext.bn_finalize(gsum, gsq, gamma, beta, rm, rv, nbt,
                                 scale, shift, m_count, float(bn.momentum),
                                 float(bn.eps), 1 if track else 0)

            return torch.addcmul(shift.view(1, 16, 1, 1, 1), raw,
                                 scale.view(1, 16, 1, 1, 1))
        else:
            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, w, 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]
```

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.
