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 2 of 6 for this task. Your previous solution (turn 1) 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


# ============================================================================
# Custom CUDA kernels
#
# Key optimizations:
#  1. ConvTranspose3d(s=2,k=3,p=1) is decomposed into a "sub-pixel" direct
#     convolution: every output voxel only depends on a 2x2x2 input
#     neighborhood (parity class dependent), so no work is wasted on
#     zero-stuffed elements.
#  2. BatchNorm (eval) is folded into a per-channel scale/shift.
#  3. The two AvgPool3d(2) layers are exactly equivalent to a single
#     AvgPool3d(4, stride=4) (63 -> 31 -> 15 == 63 -> 15 with remainder
#     dropped, and mean-of-means == global mean for equal-size blocks).
#  4. EVERYTHING is fused into one kernel: each thread block computes a
#     4x4x4 conv-transpose tile for all 16 output channels, applies BN, and
#     reduces it to the final pooled value.  The 1 GB intermediate
#     conv-transpose tensor is never materialized in DRAM.
#  5. A second kernel recomputes batch statistics on the fly (training mode
#     only), so the custom path is used in both train and eval modes.
# ============================================================================

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 64
#define PITCH (CO + 1)

// For local coordinate l in [0,4) of a 4-wide output tile:
//   global D = 4*tile + 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;
}

// Fused ConvTranspose3d + BN(scale/shift) + AvgPool(4,4,4)
// x: (N,3,32,32,32)  w: (3,16,3,3,3)  out: (N,16,15,15,15)
__global__ void fused_convT_bn_pool4_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 * 27];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_par[3 * CO];
    __shared__ float s_red[NTHREADS][PITCH];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) {
        s_par[tid]          = bias[tid];
        s_par[CO + tid]     = scale[tid] * (1.0f / (float)NTHREADS);
        s_par[2 * CO + tid] = shift[tid];
    }
    __syncthreads();

    const int vd = tid >> 4;
    const int vh = (tid >> 2) & 3;
    const int vw = tid & 3;
    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);

    const int TILES = ODIM * ODIM * ODIM;

    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TILES;
        const int rem = t - n * TILES;
        const int td  = rem / (ODIM * ODIM);
        const int th  = (rem / ODIM) % ODIM;
        const int tw  = rem % ODIM;
        const int bd = 2 * td, bh = 2 * th, bw = 2 * tw;

        const float* xbase = x + (long)n * CI * IDIM * IDIM * IDIM;
        for (int i = tid; i < CI * 27; i += NTHREADS) {
            const int ic = i / 27;
            const int a  = (i / 9) % 3;
            const int b  = (i / 3) % 3;
            const int c  = i % 3;
            s_x[i] = xbase[((ic * IDIM + (bd + a)) * IDIM + (bh + b)) * IDIM + (bw + c)];
        }
        __syncthreads();

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

        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 + 27];
                    const float x2 = s_x[xs + 54];
                    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];
                    }
                }
            }
        }

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

        if (tid < CO) {
            float s = 0.0f;
            #pragma unroll 8
            for (int i = 0; i < NTHREADS; ++i) s += s_red[i][tid];
            const int oc = tid;
            const long oidx = (((long)n * CO + oc) * ODIM + td) * (ODIM * ODIM)
                            + (long)th * ODIM + tw;
            out[oidx] = s * s_par[CO + oc] + s_par[2 * CO + oc];
        }
        __syncthreads();
    }
}

// Computes per-channel sum and sum-of-squares of the full ConvTranspose3d
// output (N,16,63,63,63) without materializing it (used for training-mode
// BatchNorm statistics). Accumulated in double precision via atomics.
__global__ void convT_stats_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    double* __restrict__ gsum,
    double* __restrict__ gsq,
    int total_tiles)
{
    const int TD = 16;  // ceil(63/4) tiles per axis
    __shared__ float s_x[CI * 27];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_red2[NTHREADS][PITCH];

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

    const int vd = tid >> 4;
    const int vh = (tid >> 2) & 3;
    const int vw = tid & 3;
    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);

    const int TILES = TD * TD * TD;

    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TILES;
        const int rem = t - n * TILES;
        const int td  = rem / (TD * TD);
        const int th  = (rem / TD) % TD;
        const int tw  = rem % TD;
        const int bd = 2 * td, bh = 2 * th, bw = 2 * tw;

        const float* xbase = x + (long)n * CI * IDIM * IDIM * IDIM;
        for (int i = tid; i < CI * 27; i += NTHREADS) {
            const int ic = i / 27;
            const int a  = (i / 9) % 3;
            const int b  = (i / 3) % 3;
            const int c  = i % 3;
            const int gi = bd + a, gj = bh + b, gk = bw + c;
            float v = 0.0f;
            if (gi < IDIM && gj < IDIM && gk < IDIM)
                v = xbase[((ic * IDIM + gi) * IDIM + gj) * IDIM + gk];
            s_x[i] = v;
        }
        __syncthreads();

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

        const bool valid = (4 * td + vd) < 63 && (4 * th + vh) < 63 && (4 * tw + vw) < 63;
        if (valid) {
            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 + 27];
                        const float x2 = s_x[xs + 54];
                        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];
                        }
                    }
                }
            }
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] += bias[oc];
        }

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

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

torch::Tensor fused_forward(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 * ODIM * ODIM * ODIM;
    int blocks = total < 4096 ? total : 4096;
    fused_convT_bn_pool4_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 conv_stats(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                torch::Tensor gsum, torch::Tensor gsq) {
    const int N = x.size(0);
    const int total = N * 16 * 16 * 16;
    int blocks = total < 4096 ? total : 4096;
    convT_stats_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        gsum.data_ptr<double>(), gsq.data_ptr<double>(), total);
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor scale, torch::Tensor shift);\n"
    "void conv_stats(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor gsum, torch::Tensor gsq);"
)

_ext = load_inline(
    name="fused_convT_bn_pool",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward", "conv_stats"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: ConvTranspose3d + BatchNorm3d + two AvgPool3d
    layers are fused into a single custom CUDA kernel (inference).  In
    training mode an additional custom kernel computes the BatchNorm batch
    statistics on the fly so the fused path is used in both modes.
    """

    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
        beta = bn.bias
        eps = bn.eps

        if self.training:
            n = x.shape[0]
            m_count = float(n) * (63.0 ** 3)  # elements per channel (exactly representable)
            gsum = torch.zeros(16, dtype=torch.float64, device=x.device)
            gsq = torch.zeros(16, dtype=torch.float64, device=x.device)
            self.ext.conv_stats(x, w, b, gsum, gsq)

            mean = gsum / m_count
            var = gsq / m_count - mean * mean
            var = torch.clamp(var, min=0.0)
            inv_std = torch.rsqrt(var + eps)
            scale = (gamma.double() * inv_std).float()
            shift = (beta.double() - mean * (gamma.double() * inv_std)).float()

            if bn.track_running_stats:
                with torch.no_grad():
                    mom = bn.momentum
                    var_unb = var * (m_count / (m_count - 1.0))
                    bn.running_mean.mul_(1.0 - mom).add_(mean.float(), alpha=mom)
                    bn.running_var.mul_(1.0 - mom).add_(var_unb.float(), alpha=mom)
                    bn.num_batches_tracked.add_(1)
        else:
            scale = gamma * torch.rsqrt(bn.running_var + eps)
            shift = beta - bn.running_mean * scale

        return self.ext.fused_forward(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]
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 4.95 ms; the PyTorch reference (with TF32 enabled) takes 36.3 ms; speedup 7.333x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| convT_stats_kernel(float const*, float const*, float const*, double*, double*, int) | 3281.7 | 1 |
| fused_convT_bn_pool4_kernel(float const*, float const*, float const*, float const*, float  | 1551.0 | 1 |
| void at::native::unrolled_elementwise_kernel<at::native::direct_copy_kernel_cuda(at::Tenso | 8.2 | 4 |
| void at::native::unrolled_elementwise_kernel<at::native::direct_copy_kernel_cuda(at::Tenso | 6.7 | 3 |
| void at::native::vectorized_elementwise_kernel<2, at::native::BinaryFunctor<double, double | 6.7 | 4 |
| void at::native::vectorized_elementwise_kernel<2, at::native::BUnaryFunctor<double, double | 4.2 | 2 |
| void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<float, float,  | 3.7 | 2 |
| void at::native::vectorized_elementwise_kernel<2, at::native::FillFunctor<double>, std::ar | 3.7 | 2 |

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.
