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

template <int CT, 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[CT][81];
    __shared__ float pool_pairs[CT][128];
    __shared__ float warp_sums[CT][8];
    __shared__ float warp_squares[CT][8];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int channel_base = blockIdx.y * CT;
    const int batch = blockIdx.z;

    for (int i = tid; i < CT * 81; i += 256) {
        const int c = i / 81;
        const int j = i % 81;
        const int channel = channel_base + c;
        weights[c][j] = channel < channels
            ? weight[((j / 27) * channels + channel) * 27 + j % 27]
            : 0.0f;
    }
    __syncthreads();

    const int ix = lane;
    const int iy = 4 * (blockIdx.x & 7) + ((tid >> 5) & 3);
    const int iz = 2 * (blockIdx.x >> 3) + (tid >> 7);
    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;

    float v[CT][8];

    #pragma unroll
    for (int c = 0; c < CT; ++c) {
        const float b = channel_base + c < channels
            ? bias[channel_base + c] : 0.0f;
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            v[c][j] = b;
    }

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

        const float a000 = p[0];
        const float a010 = has_y ? p[32] : 0.0f;
        const float a100 = has_z ? p[1024] : 0.0f;
        const float a110 = has_y && has_z ? p[1056] : 0.0f;

        float a001 = __shfl_down_sync(0xffffffff, a000, 1);
        float a011 = __shfl_down_sync(0xffffffff, a010, 1);
        float a101 = __shfl_down_sync(0xffffffff, a100, 1);
        float a111 = __shfl_down_sync(0xffffffff, a110, 1);

        if (!has_x)
            a001 = a011 = a101 = a111 = 0.0f;

        #pragma unroll
        for (int c = 0; c < CT; ++c) {
            const float* k = weights[c] + ic * 27;

            v[c][0] = fmaf(a000, k[13], v[c][0]);

            v[c][1] = fmaf(a000, k[14], v[c][1]);
            v[c][1] = fmaf(a001, k[12], v[c][1]);

            v[c][2] = fmaf(a000, k[16], v[c][2]);
            v[c][2] = fmaf(a010, k[10], v[c][2]);

            v[c][3] = fmaf(a000, k[17], v[c][3]);
            v[c][3] = fmaf(a001, k[15], v[c][3]);
            v[c][3] = fmaf(a010, k[11], v[c][3]);
            v[c][3] = fmaf(a011, k[9], v[c][3]);

            v[c][4] = fmaf(a000, k[22], v[c][4]);
            v[c][4] = fmaf(a100, k[4], v[c][4]);

            v[c][5] = fmaf(a000, k[23], v[c][5]);
            v[c][5] = fmaf(a001, k[21], v[c][5]);
            v[c][5] = fmaf(a100, k[5], v[c][5]);
            v[c][5] = fmaf(a101, k[3], v[c][5]);

            v[c][6] = fmaf(a000, k[25], v[c][6]);
            v[c][6] = fmaf(a010, k[19], v[c][6]);
            v[c][6] = fmaf(a100, k[7], v[c][6]);
            v[c][6] = fmaf(a110, k[1], v[c][6]);

            v[c][7] = fmaf(a000, k[26], v[c][7]);
            v[c][7] = fmaf(a001, k[24], v[c][7]);
            v[c][7] = fmaf(a010, k[20], v[c][7]);
            v[c][7] = fmaf(a011, k[18], v[c][7]);
            v[c][7] = fmaf(a100, k[8], v[c][7]);
            v[c][7] = fmaf(a101, k[6], v[c][7]);
            v[c][7] = fmaf(a110, k[2], v[c][7]);
            v[c][7] = fmaf(a111, k[0], v[c][7]);
        }
    }

    #pragma unroll
    for (int c = 0; c < CT; ++c) {
        if (!has_x)
            v[c][1] = v[c][3] = v[c][5] = v[c][7] = 0.0f;
        if (!has_y)
            v[c][2] = v[c][3] = v[c][6] = v[c][7] = 0.0f;
        if (!has_z)
            v[c][4] = v[c][5] = v[c][6] = v[c][7] = 0.0f;

        float sum = ((v[c][0] + v[c][1]) + (v[c][2] + v[c][3]))
                  + ((v[c][4] + v[c][5]) + (v[c][6] + v[c][7]));

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

        if (TRAINING) {
            float square = v[c][0] * v[c][0];
            #pragma unroll
            for (int j = 1; j < 8; ++j)
                square = fmaf(v[c][j], v[c][j], square);

            sum = warp_sum(sum);
            square = warp_sum(square);
            if (lane == 0) {
                warp_sums[c][tid >> 5] = sum;
                warp_squares[c][tid >> 5] = square;
            }
        }
    }

    __syncthreads();

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

        #pragma unroll
        for (int c = 0; c < CT; ++c) {
            const int channel = channel_base + c;

            const float pool_sum =
                (pool_pairs[c][base] + pool_pairs[c][base + 16]) +
                (pool_pairs[c][base + 64] + pool_pairs[c][base + 80]);

            if (channel < channels && gx < 15 && gy < 15 && gz < 15) {
                const int index =
                    (((batch * channels + channel) * 15 + gz) * 15 + gy)
                    * 15 + gx;

                float result = pool_sum * (1.0f / 64.0f);
                if (!TRAINING) {
                    const float2 a = coefficients[channel];
                    result = fmaf(result, a.x, a.y);
                }
                pooled[index] = result;
            }

            if (TRAINING) {
                float block_sum = tid < 8 ? warp_sums[c][tid] : 0.0f;
                float block_square = tid < 8 ? warp_squares[c][tid] : 0.0f;
                block_sum = warp_sum(block_sum);
                block_square = warp_sum(block_square);

                if (tid == 0 && channel < channels) {
                    partial[(channel * batch_size + batch) * 128
                            + 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 * 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 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);
    constexpr int CT = 4;

    auto output = torch::empty({batch, channels, 15, 15, 15}, x.options());
    auto coefficients = torch::empty({channels, 2}, x.options());
    const dim3 grid(128, (channels + CT - 1) / CT, batch);

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

        conv_pool_moments<CT, 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<CT, 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_channel_tiled_v4",
    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,
        )
```