```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 kernel: BatchNorm3d (eval-mode affine) + two AvgPool3d(2) layers.
# Two consecutive 2x2x2 average pools (stride 2) are mathematically identical
# to a single 4x4x4 average pool with stride 4, and since BN is a per-channel
# affine transform, avg(BN(x)) = alpha * avg(x) + beta. This kernel reads the
# conv-transpose output exactly once (the 4x4x4 windows are disjoint) and never
# materializes the BN / intermediate pool tensors.
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

__global__ void fused_bn_double_avgpool_kernel(
    const float* __restrict__ x,
    const float* __restrict__ alpha,
    const float* __restrict__ beta,
    float* __restrict__ out,
    const int C,
    const int Di, const int Hi, const int Wi,
    const int Do, const int Ho, const int Wo,
    const int total)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int w = idx % Wo;
    int t = idx / Wo;
    const int h = t % Ho; t /= Ho;
    const int d = t % Do; t /= Do;
    const int c = t % C;
    const int n = t / C;

    const float* base = x + ((size_t)n * C + c) * (size_t)Di * Hi * Wi;
    const int d0 = d * 4, h0 = h * 4, w0 = w * 4;

    float sum = 0.0f;
    #pragma unroll
    for (int a = 0; a < 4; ++a) {
        #pragma unroll
        for (int b = 0; b < 4; ++b) {
            const float* row = base + ((size_t)(d0 + a) * Hi + (h0 + b)) * Wi + w0;
            sum += row[0] + row[1] + row[2] + row[3];
        }
    }

    // alpha * (sum / 64) + beta  ==  double-avgpool(BN(x))
    out[idx] = fmaf(sum, alpha[c] * 0.015625f, beta[c]);
}

torch::Tensor fused_bn_double_avgpool(torch::Tensor x, torch::Tensor alpha, torch::Tensor beta)
{
    TORCH_CHECK(x.is_cuda(), "input must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(x.dim() == 5, "input must be NCDHW");

    auto xc = x.contiguous();
    auto a = alpha.contiguous();
    auto b = beta.contiguous();

    const int64_t N  = xc.size(0);
    const int64_t C  = xc.size(1);
    const int64_t Di = xc.size(2);
    const int64_t Hi = xc.size(3);
    const int64_t Wi = xc.size(4);

    const int64_t Do = (Di - 4) / 4 + 1;
    const int64_t Ho = (Hi - 4) / 4 + 1;
    const int64_t Wo = (Wi - 4) / 4 + 1;
    TORCH_CHECK(Do > 0 && Ho > 0 && Wo > 0, "input spatial dims too small");

    auto out = at::empty({N, C, Do, Ho, Wo}, xc.options());
    const int64_t total = N * C * Do * Ho * Wo;
    if (total == 0) return out;
    TORCH_CHECK(total < 2147483647LL, "tensor too large");

    const int block = 256;
    const int64_t grid = (total + block - 1) / block;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    fused_bn_double_avgpool_kernel<<<(int)grid, block, 0, stream>>>(
        xc.data_ptr<float>(),
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)C,
        (int)Di, (int)Hi, (int)Wi,
        (int)Do, (int)Ho, (int)Wo,
        (int)total);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_cpp_source = "torch::Tensor fused_bn_double_avgpool(torch::Tensor x, torch::Tensor alpha, torch::Tensor beta);"

_fused_ops = load_inline(
    name="fused_bn_double_avgpool_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_bn_double_avgpool"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model:
      - ConvTranspose3d is kept (dispatched to cuDNN).
      - BatchNorm3d + AvgPool3d(2) + AvgPool3d(2) are replaced by a single
        custom CUDA kernel that applies the BN affine transform fused into a
        4x4x4 stride-4 average pool (mathematically equivalent), reading the
        large conv output exactly once.
      - In training mode (or on non-CUDA / non-fp32 inputs) it falls back to
        the exact original PyTorch ops.
    """
    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._fused = _fused_ops

    def forward(self, x):
        bn = self.batch_norm
        can_fuse = (
            not self.training
            and x.is_cuda
            and x.dtype == torch.float32
            and bn.running_mean is not None
            and bn.running_var is not None
        )

        if not can_fuse:
            x = self.conv_transpose(x)
            x = self.batch_norm(x)
            x = self.avg_pool1(x)
            x = self.avg_pool2(x)
            return x

        # Per-channel eval-mode BN affine parameters.
        inv_std = 1.0 / torch.sqrt(bn.running_var + bn.eps)
        if bn.weight is not None:
            alpha = (bn.weight * inv_std).contiguous()
        else:
            alpha = inv_std.contiguous()
        beta = -bn.running_mean * alpha
        if bn.bias is not None:
            beta = beta + bn.bias
        beta = beta.contiguous()

        ct = self.conv_transpose
        y = F.conv_transpose3d(
            x, ct.weight, ct.bias,
            ct.stride, ct.padding, ct.output_padding,
            ct.groups, ct.dilation,
        )

        # Single fused kernel: BN affine + both average pools.
        return self._fused.fused_bn_double_avgpool(y, alpha, beta)
```