```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 x) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        x += __shfl_down_sync(0xffffffff, x, offset);
    return x;
}

// Each thread computes two adjacent input cells, reusing their overlapping
// input neighborhoods and each convolution coefficient.
template <bool TRAINING>
__global__ __launch_bounds__(256, 2) void conv_pool_moments(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ pooled,
    float2* __restrict__ partial,
    const float2* __restrict__ coefficients,
    int channels,
    int batch_size)
{
    __shared__ float weights[81];
    __shared__ float pool_sums[256];
    __shared__ float warp_sums[8];
    __shared__ float warp_squares[8];

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

    if (tid < 81)
        weights[tid] =
            weight[((tid / 27) * channels + channel) * 27 + tid % 27];

    __syncthreads();

    const int ix = 2 * (tid & 15);
    const int iy = 8 * (blockIdx.x & 3) + ((tid >> 4) & 7);
    const int iz = 2 * (blockIdx.x >> 2) + (tid >> 7);

    const bool next_x = ix < 30;
    const bool next_y = iy < 31;
    const bool next_z = iz < 31;
    const int spatial = (iz * 32 + iy) * 32 + ix;

    float v[2][8];
    const float b = bias[channel];

    #pragma unroll
    for (int q = 0; q < 2; ++q) {
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            v[q][j] = b;
    }

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

        const float a000 = p[0];
        const float a001 = p[1];
        const float a002 = next_x ? p[2] : 0.0f;

        const float a010 = next_y ? p[32] : 0.0f;
        const float a011 = next_y ? p[33] : 0.0f;
        const float a012 = next_y && next_x ? p[34] : 0.0f;

        const float a100 = next_z ? p[1024] : 0.0f;
        const float a101 = next_z ? p[1025] : 0.0f;
        const float a102 = next_z && next_x ? p[1026] : 0.0f;

        const float a110 = next_y && next_z ? p[1056] : 0.0f;
        const float a111 = next_y && next_z ? p[1057] : 0.0f;
        const float a112 =
            next_y && next_z && next_x ? p[1058] : 0.0f;

        #define FMA_PAIR(J, K, A, B)               \
        {                                         \
            const float kw = k[K];                \
            v[0][J] = fmaf(A, kw, v[0][J]);        \
            v[1][J] = fmaf(B, kw, v[1][J]);        \
        }

        FMA_PAIR(0, 13, a000, a001);

        FMA_PAIR(1, 14, a000, a001);
        FMA_PAIR(1, 12, a001, a002);

        FMA_PAIR(2, 16, a000, a001);
        FMA_PAIR(2, 10, a010, a011);

        FMA_PAIR(3, 17, a000, a001);
        FMA_PAIR(3, 15, a001, a002);
        FMA_PAIR(3, 11, a010, a011);
        FMA_PAIR(3,  9, a011, a012);

        FMA_PAIR(4, 22, a000, a001);
        FMA_PAIR(4,  4, a100, a101);

        FMA_PAIR(5, 23, a000, a001);
        FMA_PAIR(5, 21, a001, a002);
        FMA_PAIR(5,  5, a100, a101);
        FMA_PAIR(5,  3, a101, a102);

        FMA_PAIR(6, 25, a000, a001);
        FMA_PAIR(6, 19, a010, a011);
        FMA_PAIR(6,  7, a100, a101);
        FMA_PAIR(6,  1, a110, a111);

        FMA_PAIR(7, 26, a000, a001);
        FMA_PAIR(7, 24, a001, a002);
        FMA_PAIR(7, 20, a010, a011);
        FMA_PAIR(7, 18, a011, a012);
        FMA_PAIR(7,  8, a100, a101);
        FMA_PAIR(7,  6, a101, a102);
        FMA_PAIR(7,  2, a110, a111);
        FMA_PAIR(7,  0, a111, a112);

        #undef FMA_PAIR
    }

    if (!next_x)
        v[1][1] = v[1][3] = v[1][5] = v[1][7] = 0.0f;

    #pragma unroll
    for (int q = 0; q < 2; ++q) {
        if (!next_y)
            v[q][2] = v[q][3] = v[q][6] = v[q][7] = 0.0f;
        if (!next_z)
            v[q][4] = v[q][5] = v[q][6] = v[q][7] = 0.0f;
    }

    const float sum0 =
        ((v[0][0] + v[0][1]) + (v[0][2] + v[0][3])) +
        ((v[0][4] + v[0][5]) + (v[0][6] + v[0][7]));
    const float sum1 =
        ((v[1][0] + v[1][1]) + (v[1][2] + v[1][3])) +
        ((v[1][4] + v[1][5]) + (v[1][6] + v[1][7]));
    float sum = sum0 + sum1;
    pool_sums[tid] = sum;

    if (TRAINING) {
        float square0 = v[0][0] * v[0][0];
        float square1 = v[1][0] * v[1][0];

        #pragma unroll
        for (int j = 1; j < 8; ++j) {
            square0 = fmaf(v[0][j], v[0][j], square0);
            square1 = fmaf(v[1][j], v[1][j], square1);
        }

        float square = warp_sum(square0 + square1);
        sum = warp_sum(sum);

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

    __syncthreads();

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

        if (gx < 15 && gy < 15 && gz < 15) {
            float result =
                ((pool_sums[base] + pool_sums[base + 16]) +
                 (pool_sums[base + 128] + pool_sums[base + 144]))
                * (1.0f / 64.0f);

            if (!TRAINING) {
                const float2 a = coefficients[channel];
                result = fmaf(result, a.x, a.y);
            }

            const int index =
                (((batch * channels + channel) * 15 + gz) * 15 + gy)
                * 15 + gx;
            pooled[index] = result;
        }
    }

    if (TRAINING && tid < 32) {
        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) * 64 + 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 * 64;

    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 scale =
                weight[channel] * rsqrtf(float(var_d) + 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(64, channels, batch);

    if (training) {
        auto partial = torch::empty({channels, batch * 64, 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>()),
            nullptr, 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));

        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);
    } else {
        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));

        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,
            reinterpret_cast<const float2*>(coefficients.data_ptr<float>()),
            channels, batch);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_ops = load_inline(
    name="direct_conv3d_bn_double_pool_spatial_pair_v5",
    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 x.shape[0] * 3 * 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 not None
                and 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,
        )
```