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 4 of 6 for this task. Your previous solution (turn 3) was:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


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

torch::Tensor conv_bn_pool_cuda(
    torch::Tensor x, torch::Tensor conv_weight, torch::Tensor conv_bias,
    torch::Tensor bn_weight, torch::Tensor bn_bias,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum, bool training);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

template <typename T>
__device__ __forceinline__ T warp_sum(T value) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        value += __shfl_down_sync(0xffffffff, value, offset);
    return value;
}

template <bool TRAINING>
__global__ void conv_pool_moments(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ pooled,
    float2* __restrict__ partial,
    int channels,
    int batch_size)
{
    __shared__ float w[81];
    __shared__ float pool_pairs[128];
    __shared__ float warp_sums[8];
    __shared__ float warp_squares[8];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int channel = blockIdx.y;
    const int batch = blockIdx.z;

    if (tid < 81) {
        const int ic = tid / 27;
        const int k = tid % 27;
        w[tid] = weight[(ic * channels + channel) * 27 + k];
    }
    __syncthreads();

    // Each warp accesses 32 consecutive input elements.
    // A block covers a 32 x 4 x 2 input tile.
    const int ix = lane;
    const int iy = 4 * (blockIdx.x & 7) + ((tid >> 5) & 3);
    const int iz = 2 * (blockIdx.x >> 3) + (tid >> 7);

    const bool has_x = ix < 31;
    const bool has_y = iy < 31;
    const bool has_z = iz < 31;

    const int spatial_index = (iz * 32 + iy) * 32 + ix;
    const float b = bias[channel];

    float v000 = b, v001 = b, v010 = b, v011 = b;
    float v100 = b, v101 = b, v110 = b, v111 = b;

    #pragma unroll
    for (int ic = 0; ic < 3; ++ic) {
        const float* p = x + (batch * 3 + ic) * 32768 + spatial_index;
        const float* k = w + ic * 27;

        const float a000 = p[0];
        const float a001 = has_x ? p[1] : 0.0f;
        const float a010 = has_y ? p[32] : 0.0f;
        const float a011 = has_x && has_y ? p[33] : 0.0f;
        const float a100 = has_z ? p[1024] : 0.0f;
        const float a101 = has_x && has_z ? p[1025] : 0.0f;
        const float a110 = has_y && has_z ? p[1056] : 0.0f;
        const float a111 = has_x && has_y && has_z ? p[1057] : 0.0f;

        v000 = fmaf(a000, k[13], v000);

        v001 = fmaf(a000, k[14], v001);
        v001 = fmaf(a001, k[12], v001);

        v010 = fmaf(a000, k[16], v010);
        v010 = fmaf(a010, k[10], v010);

        v011 = fmaf(a000, k[17], v011);
        v011 = fmaf(a001, k[15], v011);
        v011 = fmaf(a010, k[11], v011);
        v011 = fmaf(a011, k[9], v011);

        v100 = fmaf(a000, k[22], v100);
        v100 = fmaf(a100, k[4], v100);

        v101 = fmaf(a000, k[23], v101);
        v101 = fmaf(a001, k[21], v101);
        v101 = fmaf(a100, k[5], v101);
        v101 = fmaf(a101, k[3], v101);

        v110 = fmaf(a000, k[25], v110);
        v110 = fmaf(a010, k[19], v110);
        v110 = fmaf(a100, k[7], v110);
        v110 = fmaf(a110, k[1], v110);

        v111 = fmaf(a000, k[26], v111);
        v111 = fmaf(a001, k[24], v111);
        v111 = fmaf(a010, k[20], v111);
        v111 = fmaf(a011, k[18], v111);
        v111 = fmaf(a100, k[8], v111);
        v111 = fmaf(a101, k[6], v111);
        v111 = fmaf(a110, k[2], v111);
        v111 = fmaf(a111, k[0], v111);
    }

    if (!has_x)
        v001 = v011 = v101 = v111 = 0.0f;
    if (!has_y)
        v010 = v011 = v110 = v111 = 0.0f;
    if (!has_z)
        v100 = v101 = v110 = v111 = 0.0f;

    float sum = ((v000 + v001) + (v010 + v011))
              + ((v100 + v101) + (v110 + v111));

    const float pair_sum =
        sum + __shfl_xor_sync(0xffffffff, sum, 1);
    if ((tid & 1) == 0)
        pool_pairs[tid >> 1] = pair_sum;

    if (TRAINING) {
        float square = v000 * v000;
        square = fmaf(v001, v001, square);
        square = fmaf(v010, v010, square);
        square = fmaf(v011, v011, square);
        square = fmaf(v100, v100, square);
        square = fmaf(v101, v101, square);
        square = fmaf(v110, v110, square);
        square = fmaf(v111, v111, square);

        sum = warp_sum(sum);
        square = warp_sum(square);
        if (lane == 0) {
            warp_sums[tid >> 5] = sum;
            warp_squares[tid >> 5] = square;
        }
    }

    __syncthreads();

    if (tid < 32) {
        const int gx = tid & 15;
        const int gy = 2 * (blockIdx.x & 7) + (tid >> 4);
        const int gz = blockIdx.x >> 3;
        const int base = (tid >> 4) * 32 + gx;

        const float pool_sum =
            (pool_pairs[base] + pool_pairs[base + 16]) +
            (pool_pairs[base + 64] + pool_pairs[base + 80]);

        if (gx < 15 && gy < 15 && gz < 15) {
            const int index =
                (((batch * channels + channel) * 15 + gz) * 15 + gy)
                * 15 + gx;
            pooled[index] = pool_sum * (1.0f / 64.0f);
        }

        if (TRAINING) {
            float block_sum = tid < 8 ? warp_sums[tid] : 0.0f;
            float block_square = tid < 8 ? warp_squares[tid] : 0.0f;
            block_sum = warp_sum(block_sum);
            block_square = warp_sum(block_square);
            if (tid == 0) {
                partial[(channel * batch_size + batch) * 128 + blockIdx.x] =
                    make_float2(block_sum, block_square);
            }
        }
    }
}

__global__ void training_coefficients(
    const float2* __restrict__ partial,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float2* __restrict__ coefficients,
    int batch_size,
    float eps,
    float momentum)
{
    __shared__ double warp_sums[8];
    __shared__ double warp_squares[8];

    const int channel = blockIdx.x;
    const int tid = threadIdx.x;
    const int chunks = batch_size * 128;

    double sum = 0.0;
    double square = 0.0;
    for (int i = tid; i < chunks; i += 256) {
        const float2 value = partial[channel * chunks + i];
        sum += double(value.x);
        square += double(value.y);
    }

    sum = warp_sum(sum);
    square = warp_sum(square);
    if ((tid & 31) == 0) {
        warp_sums[tid >> 5] = sum;
        warp_squares[tid >> 5] = square;
    }
    __syncthreads();

    if (tid < 32) {
        sum = tid < 8 ? warp_sums[tid] : 0.0;
        square = tid < 8 ? warp_squares[tid] : 0.0;
        sum = warp_sum(sum);
        square = warp_sum(square);

        if (tid == 0) {
            const double count = double(batch_size) * 250047.0;
            const double mean_d = sum / count;
            const double var_d = fmax(square / count - mean_d * mean_d, 0.0);
            const float mean = float(mean_d);
            const float variance = float(var_d);
            const float scale = weight[channel] * rsqrtf(variance + eps);

            coefficients[channel] =
                make_float2(scale, bias[channel] - mean * scale);

            running_mean[channel] =
                (1.0f - momentum) * running_mean[channel] + momentum * mean;
            running_var[channel] =
                (1.0f - momentum) * running_var[channel]
                + momentum * float(var_d * count / (count - 1.0));
        }
    }
}

__global__ void evaluation_coefficients(
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float2* __restrict__ coefficients,
    int channels,
    float eps)
{
    const int channel = blockIdx.x * blockDim.x + threadIdx.x;
    if (channel < channels) {
        const float scale = weight[channel] * rsqrtf(running_var[channel] + eps);
        coefficients[channel] =
            make_float2(scale, bias[channel] - running_mean[channel] * scale);
    }
}

__global__ void apply_affine(
    float* __restrict__ output,
    const float2* __restrict__ coefficients,
    int elements,
    int channels)
{
    const int vector_index = blockIdx.x * blockDim.x + threadIdx.x;
    const int index = vector_index * 4;

    if (index + 3 < elements) {
        float4 value = reinterpret_cast<const float4*>(output)[vector_index];
        const int plane = index / 3375;
        const float2 a = coefficients[plane % channels];

        if (index + 3 < (plane + 1) * 3375) {
            value.x = fmaf(value.x, a.x, a.y);
            value.y = fmaf(value.y, a.x, a.y);
            value.z = fmaf(value.z, a.x, a.y);
            value.w = fmaf(value.w, a.x, a.y);
        } else {
            const float2 a1 = coefficients[((index + 1) / 3375) % channels];
            const float2 a2 = coefficients[((index + 2) / 3375) % channels];
            const float2 a3 = coefficients[((index + 3) / 3375) % channels];
            value.x = fmaf(value.x, a.x, a.y);
            value.y = fmaf(value.y, a1.x, a1.y);
            value.z = fmaf(value.z, a2.x, a2.y);
            value.w = fmaf(value.w, a3.x, a3.y);
        }
        reinterpret_cast<float4*>(output)[vector_index] = value;
    } else {
        for (int i = index; i < elements; ++i) {
            const float2 a = coefficients[(i / 3375) % channels];
            output[i] = fmaf(output[i], a.x, a.y);
        }
    }
}

torch::Tensor conv_bn_pool_cuda(
    torch::Tensor x, torch::Tensor conv_weight, torch::Tensor conv_bias,
    torch::Tensor bn_weight, torch::Tensor bn_bias,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum, bool training)
{
    const c10::cuda::CUDAGuard guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int batch = x.size(0);
    const int channels = conv_weight.size(1);

    auto output = torch::empty({batch, channels, 15, 15, 15}, x.options());
    auto coefficients = torch::empty({channels, 2}, x.options());
    const dim3 grid(128, channels, batch);

    if (training) {
        auto partial = torch::empty({channels, batch * 128, 2}, x.options());

        conv_pool_moments<true><<<grid, 256, 0, stream>>>(
            x.data_ptr<float>(),
            conv_weight.data_ptr<float>(),
            conv_bias.data_ptr<float>(),
            output.data_ptr<float>(),
            reinterpret_cast<float2*>(partial.data_ptr<float>()),
            channels, batch);

        training_coefficients<<<channels, 256, 0, stream>>>(
            reinterpret_cast<const float2*>(partial.data_ptr<float>()),
            bn_weight.data_ptr<float>(),
            bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            reinterpret_cast<float2*>(coefficients.data_ptr<float>()),
            batch, float(eps), float(momentum));
    } else {
        conv_pool_moments<false><<<grid, 256, 0, stream>>>(
            x.data_ptr<float>(),
            conv_weight.data_ptr<float>(),
            conv_bias.data_ptr<float>(),
            output.data_ptr<float>(),
            nullptr, channels, batch);

        evaluation_coefficients<<<(channels + 255) / 256, 256, 0, stream>>>(
            bn_weight.data_ptr<float>(),
            bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            reinterpret_cast<float2*>(coefficients.data_ptr<float>()),
            channels, float(eps));
    }

    const int elements = output.numel();
    apply_affine<<<(elements + 1023) / 1024, 256, 0, stream>>>(
        output.data_ptr<float>(),
        reinterpret_cast<const float2*>(coefficients.data_ptr<float>()),
        elements, channels);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_ops = load_inline(
    name="direct_conv3d_bn_double_pool_coalesced_v3",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["conv_bn_pool_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _ConvBNPool(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, conv_weight, conv_bias, bn_weight, bn_bias,
                running_mean, running_var, eps, momentum, training):
        ctx.eps = eps
        ctx.training_mode = training
        if training:
            ctx.save_for_backward(
                x, conv_weight, conv_bias, bn_weight, bn_bias
            )
        else:
            ctx.save_for_backward(
                x, conv_weight, conv_bias, bn_weight, bn_bias,
                running_mean.clone(), running_var.clone()
            )

        return _ops.conv_bn_pool_cuda(
            x, conv_weight, conv_bias, bn_weight, bn_bias,
            running_mean, running_var, eps, momentum, training
        )

    @staticmethod
    def backward(ctx, grad_output):
        saved = ctx.saved_tensors
        with torch.enable_grad():
            inputs = [
                tensor.detach().requires_grad_(True)
                for tensor in saved[:5]
            ]
            x, conv_weight, conv_bias, bn_weight, bn_bias = inputs
            y = F.conv_transpose3d(
                x, conv_weight, conv_bias, stride=2, padding=1
            )
            y = F.batch_norm(
                y,
                None if ctx.training_mode else saved[5],
                None if ctx.training_mode else saved[6],
                bn_weight,
                bn_bias,
                training=ctx.training_mode,
                momentum=0.0,
                eps=ctx.eps,
            )
            y = F.avg_pool3d(F.avg_pool3d(y, 2), 2)
            grads = torch.autograd.grad(y, inputs, grad_output)

        return (*grads, None, None, None, None, None)


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride, padding, bias_shape):
        super().__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):
        conv = self.conv_transpose
        bn = self.batch_norm

        specialized = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 5
            and x.shape[1] == 3
            and tuple(x.shape[2:]) == (32, 32, 32)
            and 0 < x.shape[0] <= 65535
            and 0 < conv.out_channels <= 65535
            and x.shape[0] * conv.out_channels * 32768 < 2147483647
            and conv.in_channels == 3
            and conv.kernel_size == (3, 3, 3)
            and conv.stride == (2, 2, 2)
            and conv.padding == (1, 1, 1)
            and conv.output_padding == (0, 0, 0)
            and conv.dilation == (1, 1, 1)
            and conv.groups == 1
            and conv.bias is not None
            and bn.affine
            and bn.track_running_stats
            and not torch.is_autocast_enabled()
        )

        if specialized:
            specialized = all(
                tensor.is_cuda
                and tensor.device == x.device
                and tensor.dtype == torch.float32
                and tensor.is_contiguous()
                for tensor in (
                    conv.weight, conv.bias, bn.weight, bn.bias,
                    bn.running_mean, bn.running_var
                )
            )

        if not specialized:
            return self.avg_pool2(self.avg_pool1(bn(conv(x))))

        momentum = bn.momentum
        if bn.training:
            bn.num_batches_tracked.add_(1)
            if momentum is None:
                momentum = 1.0 / bn.num_batches_tracked.item()

        return _ConvBNPool.apply(
            x.contiguous(),
            conv.weight,
            conv.bias,
            bn.weight,
            bn.bias,
            bn.running_mean,
            bn.running_var,
            bn.eps,
            0.0 if momentum is None else momentum,
            bn.training,
        )
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 0.748 ms; the PyTorch reference (with TF32 enabled) takes 36.3 ms; speedup 48.529x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| void conv_pool_moments<true>(float const*, float const*, float const*, float*, float2*, in | 660.3 | 1 |
| apply_affine(float*, float2 const*, int, int) | 12.2 | 1 |
| training_coefficients(float2 const*, float const*, float const*, float*, float*, float2*,  | 6.9 | 1 |
| void at::native::vectorized_elementwise_kernel<2, at::native::CUDAFunctorOnSelf_add<long>, | 2.1 | 1 |

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.
