```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; offset >>= 1)
        value += __shfl_down_sync(0xffffffff, value, offset);
    return value;
}

// Each thread computes a 2x2x2 group of transposed-convolution outputs.
// Groups of eight threads directly produce one final pooled output.
// The full 63x63x63 convolution result is never materialized.
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 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) {
        const int ic = tid / 27;
        const int k = tid % 27;
        w[tid] = weight[(ic * channels + channel) * 27 + k];
    }
    __syncthreads();

    const int group = blockIdx.x * 32 + (tid >> 3);
    const int gx = group & 15;
    const int gy = (group >> 4) & 15;
    const int gz = group >> 8;

    const int ix = 2 * gx + (tid & 1);
    const int iy = 2 * gy + ((tid >> 1) & 1);
    const int iz = 2 * gz + ((tid >> 2) & 1);

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

    // Odd output coordinates past index 62 do not exist.
    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));

    float pool_sum = sum;
    pool_sum += __shfl_xor_sync(0xffffffff, pool_sum, 1);
    pool_sum += __shfl_xor_sync(0xffffffff, pool_sum, 2);
    pool_sum += __shfl_xor_sync(0xffffffff, pool_sum, 4);

    if ((tid & 7) == 0 && 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 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 ((tid & 31) == 0) {
            warp_sums[tid >> 5] = sum;
            warp_squares[tid >> 5] = square;
        }
        __syncthreads();

        if (tid < 32) {
            sum = tid < 8 ? warp_sums[tid] : 0.0f;
            square = tid < 8 ? warp_squares[tid] : 0.0f;
            sum = warp_sum(sum);
            square = warp_sum(square);
            if (tid == 0) {
                partial[(channel * batch_size + batch) * 128 + blockIdx.x] =
                    make_float2(sum, 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 index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index < elements) {
        const int channel = (index / 3375) % channels;
        const float2 affine = coefficients[channel];
        output[index] = fmaf(output[index], affine.x, affine.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 + 255) / 256, 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_v2",
    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 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,
        )
```