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

// Process consecutive z planes sequentially to reuse the block's weights
// while keeping only eight convolution accumulators live.
__global__ __launch_bounds__(256, 4) 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 pool_pairs[128];
    __shared__ float warp_sums[8];
    __shared__ float warp_squares[8];

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

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

    const int ix = lane;
    const int iy = 8 * (blockIdx.x & 3) + (tid >> 5);
    const int zbase = 2 * (blockIdx.x >> 2);
    const bool has_x = ix < 31;
    const bool has_y = iy < 31;
    const float b = bias[channel];

    float sum = 0.0f;
    float square = 0.0f;

    #pragma unroll 1
    for (int dz = 0; dz < 2; ++dz) {
        const int iz = zbase + dz;
        const bool has_z = iz < 31;
        const int spatial = (iz * 32 + iy) * 32 + ix;

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

        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;

        sum += ((v000 + v001) + (v010 + v011))
             + ((v100 + v101) + (v110 + v111));

        float s = v000 * v000;
        s = fmaf(v001, v001, s);
        s = fmaf(v010, v010, s);
        s = fmaf(v011, v011, s);
        s = fmaf(v100, v100, s);
        s = fmaf(v101, v101, s);
        s = fmaf(v110, v110, s);
        s = fmaf(v111, v111, s);
        square += s;
    }

    const float pair = sum + __shfl_xor_sync(0xffffffff, sum, 1);
    if ((tid & 1) == 0)
        pool_pairs[tid >> 1] = pair;

    sum = warp_sum(sum);
    square = warp_sum(square);

    if (lane == 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) {
            const int index =
                (((batch * channels + channel) * 15 + gz) * 15 + gy)
                * 15 + gx;
            pooled[index] =
                (pool_pairs[base] + pool_pairs[base + 16]) * (1.0f / 64.0f);
        }
    }

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

// In evaluation mode, combine the convolution, both pools, and BN into
// one stride-two convolution on the original input.
__global__ void prepare_evaluation(
    const float* __restrict__ weight,
    const float* __restrict__ conv_bias,
    const float* __restrict__ bn_weight,
    const float* __restrict__ bn_bias,
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float* __restrict__ packed,
    int channels,
    float eps)
{
    const int channel = blockIdx.x;
    const int tid = threadIdx.x;
    const float scale =
        bn_weight[channel] * rsqrtf(running_var[channel] + eps);

    if (tid < 81) {
        const int ic = tid / 27;
        const int index = tid % 27;
        const int dx = index % 3;
        const int dy = (index / 3) % 3;
        const int dz = index / 9;

        const int xlo = dx == 0 ? 1 : 0;
        const int ylo = dy == 0 ? 1 : 0;
        const int zlo = dz == 0 ? 1 : 0;
        const int xhi = dx == 2 ? 0 : 2;
        const int yhi = dy == 2 ? 0 : 2;
        const int zhi = dz == 2 ? 0 : 2;

        const float* k = weight + (ic * channels + channel) * 27;
        float value = 0.0f;
        for (int z = zlo; z <= zhi; ++z)
            for (int y = ylo; y <= yhi; ++y)
                for (int x = xlo; x <= xhi; ++x)
                    value += k[(z * 3 + y) * 3 + x];

        packed[channel * 82 + tid] = value * (scale / 64.0f);
    }

    if (tid == 0)
        packed[channel * 82 + 81] =
            (conv_bias[channel] - running_mean[channel]) * scale
            + bn_bias[channel];
}

__global__ __launch_bounds__(256) void evaluation_pool(
    const float* __restrict__ x,
    const float* __restrict__ packed,
    float* __restrict__ output,
    int channels)
{
    __shared__ float k[82];

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

    if (tid < 82)
        k[tid] = packed[channel * 82 + tid];
    __syncthreads();

    const int index = blockIdx.x * 256 + tid;
    if (index >= 3375)
        return;

    const int ox = index % 15;
    const int oy = (index / 15) % 15;
    const int oz = index / 225;
    const int spatial = ((2 * oz) * 32 + 2 * oy) * 32 + 2 * ox;

    float value = k[81];
    #pragma unroll
    for (int ic = 0; ic < 3; ++ic) {
        const float* p = x + (batch * 3 + ic) * 32768 + spatial;
        #pragma unroll
        for (int dz = 0; dz < 3; ++dz) {
            #pragma unroll
            for (int dy = 0; dy < 3; ++dy) {
                #pragma unroll
                for (int dx = 0; dx < 3; ++dx) {
                    value = fmaf(
                        p[dz * 1024 + dy * 32 + dx],
                        k[ic * 27 + dz * 9 + dy * 3 + dx],
                        value);
                }
            }
        }
    }

    output[(batch * channels + channel) * 3375 + index] = value;
}

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

    if (training) {
        auto partial = torch::empty({channels, batch * 64, 2}, x.options());
        auto coefficients = torch::empty({channels, 2}, x.options());

        conv_pool_moments<<<dim3(64, channels, batch), 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));

        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 {
        auto packed = torch::empty({channels, 82}, x.options());

        prepare_evaluation<<<channels, 128, 0, stream>>>(
            conv_weight.data_ptr<float>(),
            conv_bias.data_ptr<float>(),
            bn_weight.data_ptr<float>(),
            bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            packed.data_ptr<float>(),
            channels, float(eps));

        evaluation_pool<<<dim3(14, channels, batch), 256, 0, stream>>>(
            x.data_ptr<float>(),
            packed.data_ptr<float>(),
            output.data_ptr<float>(),
            channels);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_ops = load_inline(
    name="direct_conv3d_bn_pool_sequential_z_v6",
    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,
        )
```