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):
    """
    Model that performs a 3D transposed convolution, scales the output, applies batch normalization, 
    and then performs global average pooling. 
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 16, 32, 32
kernel_size = 5
scale_factor = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scale_factor]
    ```
    
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
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA kernels.
#
# Key insight: the network ends in a *global* average pool, and BatchNorm is a
# per-channel affine transform. Therefore the final result only depends on the
# per-channel spatial sums (and sums of squares, for train-mode batch stats)
# of the transposed-convolution output -- the full output tensor never needs
# to be materialized or re-read.
#
# In eval mode (fixed running stats) we can go one step further: the spatial
# sum of a stride-1/pad-0 ConvTranspose3d output decomposes exactly as
#     sum(out[b,o]) = sum_i sum(x[b,i]) * sum(w[i,o]) + bias[o] * S_out
# so the entire 500+ GFLOP convolution collapses into two cheap reductions
# plus a tiny 16x64x128 contraction, fused with scale + BN + pooling.
#
# In training mode we run the (unavoidable) convolution and fuse
# scale + batch-stat computation + running-stat update + BN + global pooling
# into a pair of lightweight reduction kernels.
# ---------------------------------------------------------------------------

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>

#define BLOCK_SIZE 256

// Per-(batch, channel) spatial sums (and optionally sums of squares) of a
// contiguous NCDHW tensor. One block per (b, c) pair.
__global__ void spatial_sums_kernel(const float* __restrict__ x,
                                    float* __restrict__ sums,
                                    float* __restrict__ sumsq,
                                    int S, int compute_sq) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    long base = ((long)b * (long)gridDim.x + (long)c) * (long)S;
    const float* xp = x + base;

    float acc = 0.f, acc2 = 0.f;
    if ((S & 3) == 0) {
        const float4* xp4 = reinterpret_cast<const float4*>(xp);
        int S4 = S >> 2;
        for (int s = threadIdx.x; s < S4; s += BLOCK_SIZE) {
            float4 v = xp4[s];
            acc += (v.x + v.y) + (v.z + v.w);
            if (compute_sq) acc2 += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
        }
    } else {
        for (int s = threadIdx.x; s < S; s += BLOCK_SIZE) {
            float v = xp[s];
            acc += v;
            if (compute_sq) acc2 += v * v;
        }
    }

    // warp reduction
    for (int off = 16; off > 0; off >>= 1) {
        acc += __shfl_down_sync(0xffffffffu, acc, off);
        if (compute_sq) acc2 += __shfl_down_sync(0xffffffffu, acc2, off);
    }
    __shared__ float sh[BLOCK_SIZE / 32];
    __shared__ float sh2[BLOCK_SIZE / 32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh[wid] = acc;
        if (compute_sq) sh2[wid] = acc2;
    }
    __syncthreads();
    if (wid == 0) {
        int nw = BLOCK_SIZE / 32;
        float v  = (lane < nw) ? sh[lane] : 0.f;
        float v2 = (lane < nw && compute_sq) ? sh2[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
            if (compute_sq) v2 += __shfl_down_sync(0xffffffffu, v2, off);
        }
        if (lane == 0) {
            long idx = (long)b * (long)gridDim.x + (long)c;
            sums[idx] = v;
            if (compute_sq) sumsq[idx] = v2;
        }
    }
}

std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    int64_t B = xc.size(0);
    int64_t C = xc.size(1);
    int64_t S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    TORCH_CHECK(B <= 65535, "batch too large");
    auto sums  = torch::empty({B, C}, xc.options());
    auto sumsq = compute_sq ? torch::empty({B, C}, xc.options())
                            : torch::empty({0}, xc.options());
    if (B * C == 0) return {sums, sumsq};
    dim3 grid((unsigned)C, (unsigned)B);
    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sums_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), sums.data_ptr<float>(),
        compute_sq ? sumsq.data_ptr<float>() : nullptr,
        (int)S, compute_sq ? 1 : 0);
    return {sums, sumsq};
}

// Eval-mode fused kernel: contracts input sums with kernel sums, applies
// conv bias, scale, BatchNorm (running stats) and emits the globally pooled
// result. Sw is [In, Out] (ConvTranspose weight layout is (In, Out, kD,kH,kW)).
__global__ void eval_finalize_kernel(const float* __restrict__ Sx,
                                     const float* __restrict__ Sw,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ rmean,
                                     const float* __restrict__ rvar,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     float* __restrict__ out,
                                     int B, int In, int Out,
                                     float scale, float inv_Sout, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * Out) return;
    int b = idx / Out;
    int o = idx - b * Out;
    const float* sx = Sx + (long)b * In;
    float acc = 0.f;
    for (int i = 0; i < In; ++i) acc += sx[i] * Sw[(long)i * Out + o];
    float mean = acc * inv_Sout;
    if (bias != nullptr) mean += bias[o];
    mean *= scale;
    float g  = (gamma != nullptr) ? gamma[o] : 1.f;
    float be = (beta  != nullptr) ? beta[o]  : 0.f;
    out[idx] = (mean - rmean[o]) * rsqrtf(rvar[o] + eps) * g + be;
}

torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta,
                            double scale, double inv_Sout, double eps) {
    int B = Sx.size(0), In = Sx.size(1), Out = Sw.size(1);
    auto out = torch::empty({B, Out}, Sx.options());
    int total = B * Out;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    eval_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, In, Out,
        (float)scale, (float)inv_Sout, (float)eps);
    return out;
}

// Train-mode (or generic) fused kernel: given per-(b,c) sums and sums of
// squares of the convolution output, compute batch statistics, update the
// running statistics in place (exactly like PyTorch: biased variance for
// normalization, unbiased for the running average), then apply the affine
// transform and emit the globally pooled result. With training=0 the running
// statistics are used directly (eval fallback for non-unit conv params).
__global__ void train_finalize_kernel(const float* __restrict__ sums,
                                      const float* __restrict__ sumsq,
                                      float* __restrict__ rmean,
                                      float* __restrict__ rvar,
                                      const float* __restrict__ gamma,
                                      const float* __restrict__ beta,
                                      float* __restrict__ out,
                                      int B, int C, float S,
                                      float scale, float momentum,
                                      float eps, int training) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * C) return;
    int b = idx / C;
    int c = idx - b * C;
    float mean, invstd;
    float N = S * (float)B;
    if (training) {
        float s = 0.f, s2 = 0.f;
        for (int bb = 0; bb < B; ++bb) {
            s  += sums[(long)bb * C + c];
            s2 += sumsq[(long)bb * C + c];
        }
        s  *= scale;
        s2 *= scale * scale;
        mean = s / N;
        float var = s2 / N - mean * mean;
        var = fmaxf(var, 0.f);
        invstd = rsqrtf(var + eps);
        if (b == 0 && rmean != nullptr && momentum >= 0.f) {
            float unbiased = var * (N / (N - 1.f));
            rmean[c] = (1.f - momentum) * rmean[c] + momentum * mean;
            rvar[c]  = (1.f - momentum) * rvar[c]  + momentum * unbiased;
        }
    } else {
        mean = rmean[c];
        invstd = rsqrtf(rvar[c] + eps);
    }
    float g  = (gamma != nullptr) ? gamma[c] : 1.f;
    float be = (beta  != nullptr) ? beta[c]  : 0.f;
    float mbc = sums[(long)b * C + c] * (scale / S);
    out[idx] = (mbc - mean) * invstd * g + be;
}

torch::Tensor train_finalize(torch::Tensor sums, torch::Tensor sumsq,
                             torch::Tensor rmean, torch::Tensor rvar,
                             torch::Tensor gamma, torch::Tensor beta,
                             double scale, double S, double momentum,
                             double eps, bool training) {
    int B = sums.size(0), C = sums.size(1);
    auto out = torch::empty({B, C}, sums.options());
    int total = B * C;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    train_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        rmean.numel() ? rmean.data_ptr<float>() : nullptr,
        rvar.numel()  ? rvar.data_ptr<float>()  : nullptr,
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, C, (float)S, (float)scale,
        (float)momentum, (float)eps, training ? 1 : 0);
    return out;
}
"""

_cpp_src = r"""
#include <vector>
std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq);
torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta,
                            double scale, double inv_Sout, double eps);
torch::Tensor train_finalize(torch::Tensor sums, torch::Tensor sumsq,
                             torch::Tensor rmean, torch::Tensor rvar,
                             torch::Tensor gamma, torch::Tensor beta,
                             double scale, double S, double momentum,
                             double eps, bool training);
"""

_ext = load_inline(
    name="fused_convtranspose3d_bn_globalpool",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["spatial_sums", "eval_finalize", "train_finalize"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Keeps identical submodules (so state_dicts
    match) but replaces the operator chain with fused custom CUDA kernels.

    Eval mode: the ConvTranspose3d never runs -- its spatial sums are computed
    in closed form from per-channel input sums and kernel sums, then fused
    with scale + BatchNorm + global average pooling in a single tiny kernel.

    Train mode: the convolution runs (cuDNN) and scale + batch statistics +
    running-stat update + BatchNorm + global average pooling are fused into
    one reduction pass plus one tiny finalize kernel (no extra passes over the
    ~130 MB intermediate, which is also never re-read).
    """

    def __init__(self, in_channels, out_channels, kernel_size, scale_factor,
                 eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = float(scale_factor)
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        ct = self.conv_transpose
        # The closed-form sum decomposition requires the transposed conv to be
        # a "full" convolution (stride 1, no padding, unit dilation/groups).
        self._simple_conv = (
            ct.groups == 1
            and tuple(ct.stride) == (1, 1, 1)
            and tuple(ct.padding) == (0, 0, 0)
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
        )

    def _reference_forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

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

        # Fall back to the reference pipeline for unsupported configurations.
        if not (x.is_cuda and x.dtype == torch.float32):
            return self._reference_forward(x)
        if bn.training and bn.momentum is None:
            return self._reference_forward(x)  # cumulative moving average
        if (not bn.training) and bn.track_running_stats and (
            bn.running_mean is None or not bn.running_mean.is_cuda
        ):
            return self._reference_forward(x)

        B = x.shape[0]
        gamma = bn.weight if bn.affine else x.new_empty(0)
        beta = bn.bias if bn.affine else x.new_empty(0)

        if (not bn.training) and bn.track_running_stats and self._simple_conv:
            # ---- closed-form path: no convolution at all -------------------
            In, Out = ct.in_channels, ct.out_channels
            Sx = _ext.spatial_sums(x, False)[0]            # [B, In]
            Sw = _ext.spatial_sums(ct.weight, False)[0]    # [In, Out]
            D, H, W = x.shape[2], x.shape[3], x.shape[4]
            k = ct.kernel_size
            S_out = (D + k[0] - 1) * (H + k[1] - 1) * (W + k[2] - 1)
            bias = ct.bias if ct.bias is not None else x.new_empty(0)
            out = _ext.eval_finalize(
                Sx, Sw, bias, bn.running_mean, bn.running_var, gamma, beta,
                self.scale_factor, 1.0 / float(S_out), bn.eps,
            )
            return out.view(B, Out, 1, 1, 1)

        # ---- general path: conv + fused reduction --------------------------
        y = ct(x)
        C = y.shape[1]
        S = y.shape[2] * y.shape[3] * y.shape[4]
        sums, sumsq = _ext.spatial_sums(y, True)
        if bn.track_running_stats:
            rmean, rvar = bn.running_mean, bn.running_var
        else:
            rmean = rvar = x.new_empty(0)
        momentum = bn.momentum if bn.momentum is not None else -1.0
        if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
            bn.num_batches_tracked += 1
        out = _ext.train_finalize(
            sums, sumsq, rmean, rvar, gamma, beta,
            self.scale_factor, float(S), momentum, bn.eps, bn.training,
        )
        return out.view(B, C, 1, 1, 1)
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 11.8 ms; the PyTorch reference (with TF32 enabled) takes 12.8 ms; speedup 1.085x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| sm80_xmma_dgrad_implicit_gemm_indexed_tf32f32_tf32f32_f32_nhwckrsc_nhwc_tilesize64x64x64_s | 10891.6 | 1 |
| void at::native::elementwise_kernel<128, 2, at::native::gpu_kernel_impl_nocast<at::native: | 319.8 | 1 |
| void cudnn::engines_precompiled::nhwcToNchwKernel<float, float, float, true, false, (cudnn | 302.0 | 1 |
| spatial_sums_kernel(float const*, float*, float*, int, int) | 164.3 | 1 |
| void cudnn::engines_precompiled::nchwToNhwcKernel<float, float, float, false, true, (cudnn | 106.6 | 2 |
| train_finalize_kernel(float const*, float const*, float*, float*, float const*, float cons | 5.5 | 1 |
| void at::native::vectorized_elementwise_kernel<2, at::native::CUDAFunctorOnSelf_add<long>, | 2.4 | 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.
