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

# Fused CUDA implementation:
#   ConvTranspose3d (stride=2, pad=1, K=3) + BatchNorm3d + AvgPool3d(2) + AvgPool3d(2)
# Two consecutive 2x2x2 average pools are exactly one 4x4x4 average pool, and the
# per-channel affine BatchNorm transform commutes with the average, so the whole
# pipeline collapses into a single kernel producing the pooled output directly.
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

// One thread per final output element. Each final output averages a 4x4x4 block
// of the (virtual) ConvTranspose3d output. For stride=2/pad=1/K=3 an output coord
// o reads input coords: o even -> (k=1, i=o/2); o odd -> (k=0, i=(o+1)/2), (k=2, i=(o-1)/2).
__global__ void fused_convt_bn_pool_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ a_coef,
    const float* __restrict__ b_coef,
    float* __restrict__ out,
    const int B, const int Cin, const int D, const int H, const int W,
    const int Cout, const int Do, const int Ho, const int Wo)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = B * Cout * Do * Ho * Wo;
    if (idx >= total) return;

    const int ox = idx % Wo;
    const int oy = (idx / Wo) % Ho;
    const int oz = (idx / (Wo * Ho)) % Do;
    const int co = (idx / (Wo * Ho * Do)) % Cout;
    const int bn = idx / (Wo * Ho * Do * Cout);

    const int izb = 2 * oz;
    const int iyb = 2 * oy;
    const int ixb = 2 * ox;
    const int HW = H * W;

    const float* __restrict__ xb = x + (long long)bn * Cin * D * HW;
    const float* __restrict__ wc = w + (long long)co * 27;

    float sum = 0.0f;

    #pragma unroll
    for (int dz = 0; dz < 4; ++dz) {
        const int hz = dz >> 1;
        int kz[2], iz[2], nz;
        if ((dz & 1) == 0) { nz = 1; kz[0] = 1; iz[0] = izb + hz; }
        else { nz = 2; kz[0] = 0; iz[0] = izb + hz + 1; kz[1] = 2; iz[1] = izb + hz; }

        #pragma unroll
        for (int dy = 0; dy < 4; ++dy) {
            const int hy = dy >> 1;
            int ky[2], iy[2], ny;
            if ((dy & 1) == 0) { ny = 1; ky[0] = 1; iy[0] = iyb + hy; }
            else { ny = 2; ky[0] = 0; iy[0] = iyb + hy + 1; ky[1] = 2; iy[1] = iyb + hy; }

            #pragma unroll
            for (int dx = 0; dx < 4; ++dx) {
                const int hx = dx >> 1;
                int kx[2], ix[2], nx;
                if ((dx & 1) == 0) { nx = 1; kx[0] = 1; ix[0] = ixb + hx; }
                else { nx = 2; kx[0] = 0; ix[0] = ixb + hx + 1; kx[1] = 2; ix[1] = ixb + hx; }

                #pragma unroll
                for (int tz = 0; tz < 2; ++tz) {
                    if (tz >= nz) break;
                    const int zoff = iz[tz] * HW;
                    const int wz = kz[tz] * 9;
                    #pragma unroll
                    for (int ty = 0; ty < 2; ++ty) {
                        if (ty >= ny) break;
                        const int yzoff = zoff + iy[ty] * W;
                        const int wyz = wz + ky[ty] * 3;
                        #pragma unroll
                        for (int tx = 0; tx < 2; ++tx) {
                            if (tx >= nx) break;
                            const int xoff = yzoff + ix[tx];
                            const int woff = wyz + kx[tx];
                            for (int ci = 0; ci < Cin; ++ci) {
                                sum = fmaf(xb[(long long)ci * D * HW + xoff],
                                           wc[ci * Cout * 27 + woff], sum);
                            }
                        }
                    }
                }
            }
        }
    }

    out[idx] = a_coef[co] * sum + b_coef[co];
}

// Per-channel sum / sum-of-squares over the full (unpooled) ConvTranspose3d output.
// One block per (batch, out-channel); used for training-mode BatchNorm statistics.
__global__ void conv_stats_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ gsum,
    float* __restrict__ gsq,
    const int Cin, const int D, const int H, const int W,
    const int Cout, const int Dc, const int Hc, const int Wc)
{
    const int bn = blockIdx.x / Cout;
    const int co = blockIdx.x % Cout;
    const int HW = H * W;
    const float* __restrict__ xb = x + (long long)bn * Cin * D * HW;
    const float* __restrict__ wc = w + (long long)co * 27;

    float lsum = 0.0f, lsq = 0.0f;
    const int spatial = Dc * Hc * Wc;
    for (int s = threadIdx.x; s < spatial; s += blockDim.x) {
        const int ox = s % Wc;
        const int oy = (s / Wc) % Hc;
        const int oz = s / (Wc * Hc);

        int kz[2], iz[2], nz, ky[2], iy[2], ny, kx[2], ix[2], nx;
        if ((oz & 1) == 0) { nz = 1; kz[0] = 1; iz[0] = oz >> 1; }
        else { nz = 2; kz[0] = 0; iz[0] = (oz >> 1) + 1; kz[1] = 2; iz[1] = oz >> 1; }
        if ((oy & 1) == 0) { ny = 1; ky[0] = 1; iy[0] = oy >> 1; }
        else { ny = 2; ky[0] = 0; iy[0] = (oy >> 1) + 1; ky[1] = 2; iy[1] = oy >> 1; }
        if ((ox & 1) == 0) { nx = 1; kx[0] = 1; ix[0] = ox >> 1; }
        else { nx = 2; kx[0] = 0; ix[0] = (ox >> 1) + 1; kx[1] = 2; ix[1] = ox >> 1; }

        float v = 0.0f;
        for (int tz = 0; tz < nz; ++tz) {
            const int zoff = iz[tz] * HW;
            const int wz = kz[tz] * 9;
            for (int ty = 0; ty < ny; ++ty) {
                const int yzoff = zoff + iy[ty] * W;
                const int wyz = wz + ky[ty] * 3;
                for (int tx = 0; tx < nx; ++tx) {
                    const int xoff = yzoff + ix[tx];
                    const int woff = wyz + kx[tx];
                    for (int ci = 0; ci < Cin; ++ci) {
                        v = fmaf(xb[(long long)ci * D * HW + xoff],
                                 wc[ci * Cout * 27 + woff], v);
                    }
                }
            }
        }
        if (bias != nullptr) v += bias[co];
        lsum += v;
        lsq = fmaf(v, v, lsq);
    }

    __shared__ float ssum[32];
    __shared__ float ssq[32];
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        lsum += __shfl_down_sync(0xffffffffu, lsum, off);
        lsq  += __shfl_down_sync(0xffffffffu, lsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) { ssum[warp] = lsum; ssq[warp] = lsq; }
    __syncthreads();
    const int nwarp = (blockDim.x + 31) >> 5;
    if (warp == 0) {
        float v1 = (lane < nwarp) ? ssum[lane] : 0.0f;
        float v2 = (lane < nwarp) ? ssq[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            v1 += __shfl_down_sync(0xffffffffu, v1, off);
            v2 += __shfl_down_sync(0xffffffffu, v2, off);
        }
        if (lane == 0) { atomicAdd(&gsum[co], v1); atomicAdd(&gsq[co], v2); }
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor a, torch::Tensor bcoef) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto ac = a.contiguous();
    auto bc = bcoef.contiguous();

    const int B = xc.size(0), Cin = xc.size(1), D = xc.size(2), H = xc.size(3), W = xc.size(4);
    const int Cout = wc.size(1);
    const int Dc = 2 * D - 1, Hc = 2 * H - 1, Wc = 2 * W - 1;
    const int Do = Dc / 4, Ho = Hc / 4, Wo = Wc / 4;
    TORCH_CHECK(Do > 0 && Ho > 0 && Wo > 0, "spatial dims too small");

    auto out = torch::empty({B, Cout, Do, Ho, Wo}, xc.options());
    const long long total = (long long)B * Cout * Do * Ho * Wo;
    const int threads = 256;
    const long long blocks = (total + threads - 1) / threads;

    fused_convt_bn_pool_kernel<<<(int)blocks, threads>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), ac.data_ptr<float>(), bc.data_ptr<float>(),
        out.data_ptr<float>(), B, Cin, D, H, W, Cout, Do, Ho, Wo);
    return out;
}

void conv_stats(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                torch::Tensor gsum, torch::Tensor gsq) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int B = xc.size(0), Cin = xc.size(1), D = xc.size(2), H = xc.size(3), W = xc.size(4);
    const int Cout = wc.size(1);
    const int Dc = 2 * D - 1, Hc = 2 * H - 1, Wc = 2 * W - 1;

    const float* bias_ptr = (bias.numel() > 0) ? bias.contiguous().data_ptr<float>() : nullptr;
    conv_stats_kernel<<<B * Cout, 256>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
        gsum.data_ptr<float>(), gsq.data_ptr<float>(),
        Cin, D, H, W, Cout, Dc, Hc, Wc);
}
"""

cpp_source = """
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor a, torch::Tensor bcoef);
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_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward", "conv_stats"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model: ConvTranspose3d + BatchNorm3d + AvgPool3d(2) + AvgPool3d(2)
    fused into a single custom CUDA kernel (the two 2x average pools are combined
    into one 4x average pool and the affine BatchNorm is folded into the pool).
    Falls back to the native PyTorch pipeline for unsupported configurations.
    """

    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)

        ct = self.conv_transpose
        self._fast = (
            ct.kernel_size == (3, 3, 3)
            and ct.stride == (2, 2, 2)
            and ct.padding == (1, 1, 1)
            and ct.dilation == (1, 1, 1)
            and ct.output_padding == (0, 0, 0)
            and ct.groups == 1
            and self.batch_norm.affine
            and self.batch_norm.track_running_stats
        )

    def _reference(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 _coeffs(self, scale, shift):
        ct = self.conv_transpose
        a = scale * (1.0 / 64.0)
        if ct.bias is not None:
            b = shift + scale * ct.bias
        else:
            b = shift
        return a.contiguous(), b.contiguous()

    def forward(self, x):
        if (not self._fast) or (not x.is_cuda) or x.dtype != torch.float32:
            return self._reference(x)

        bn = self.batch_norm
        ct = self.conv_transpose
        x = x.contiguous()
        B, Cin, D, H, W = x.shape
        Dc, Hc, Wc = 2 * D - 1, 2 * H - 1, 2 * W - 1
        if Dc < 4 or Hc < 4 or Wc < 4:
            return self._reference(x)
        w = ct.weight.contiguous()

        if not self.training:
            scale = bn.weight / torch.sqrt(bn.running_var + bn.eps)
            shift = bn.bias - bn.running_mean * scale
            a, b = self._coeffs(scale, shift)
            return _ext.fused_forward(x, w, a, b)

        # Training mode: compute batch statistics of the full transposed-conv
        # output with a dedicated reduction kernel, then normalize with them
        # (exactly like BatchNorm3d) and refresh the running statistics.
        Cout = ct.out_channels
        gsum = torch.zeros(Cout, device=x.device, dtype=torch.float32)
        gsq = torch.zeros(Cout, device=x.device, dtype=torch.float32)
        bias_arg = ct.bias if ct.bias is not None else torch.empty(0, device=x.device)
        _ext.conv_stats(x, w, bias_arg, gsum, gsq)

        n = float(B * Dc * Hc * Wc)
        mean = gsum / n
        var_b = (gsq / n - mean * mean).clamp_(min=0.0)

        scale = bn.weight / torch.sqrt(var_b + bn.eps)
        shift = bn.bias - mean * scale
        a, b = self._coeffs(scale, shift)
        out = _ext.fused_forward(x, w, a, b)

        m = bn.momentum
        if m is None:
            bn.num_batches_tracked += 1
            m = 1.0 / float(bn.num_batches_tracked.item())
        else:
            bn.num_batches_tracked += 1
        with torch.no_grad():
            bn.running_mean.mul_(1.0 - m).add_(mean, alpha=m)
            var_u = var_b * (n / (n - 1.0))
            bn.running_var.mul_(1.0 - m).add_(var_u, alpha=m)

        return out


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