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

# ---------------------------------------------------------------------------
# Key observation (algorithmic change + fusion):
# For a transposed convolution with padding=0 (any stride/dilation), every
# input pixel contributes to ALL kH*kW kernel positions of the output. Hence
#     sum_{p,q} convT(x)[b,o,p,q] = sum_c ( sum_{i,j} x[b,c,i,j] )
#                                   * ( sum_{u,v} w[c,o,u,v] )
# so the global-average-pooled result is a tiny (B,C)@(C,O) contraction plus
# the (constant) conv bias. We never materialize the (huge) conv output.
#
# Kernel 1 computes the per-(b,c) spatial sums of the input.
# Kernel 2 fuses: contraction with per-kernel weight sums + conv bias +
#                 extra bias + log-sum-exp over channels + final scaling.
# ---------------------------------------------------------------------------

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cmath>

#define FULL_MASK 0xffffffffu

// out[b, c] = sum_{h,w} x[b, c, h, w]
__global__ void spatial_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   long long HW) {
    const long long bc = blockIdx.x;
    const float* base = x + bc * HW;

    double acc = 0.0;

    if ((reinterpret_cast<unsigned long long>(base) & 15ULL) == 0ULL) {
        const long long vec = HW >> 2;
        const float4* vbase = reinterpret_cast<const float4*>(base);
        for (long long i = threadIdx.x; i < vec; i += blockDim.x) {
            float4 v = vbase[i];
            acc += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        }
        for (long long i = (vec << 2) + threadIdx.x; i < HW; i += blockDim.x) {
            acc += (double)base[i];
        }
    } else {
        for (long long i = threadIdx.x; i < HW; i += blockDim.x) {
            acc += (double)base[i];
        }
    }

    for (int offset = 16; offset > 0; offset >>= 1)
        acc += __shfl_down_sync(FULL_MASK, acc, offset);

    __shared__ double warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = acc;
    __syncthreads();

    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        double v = (lane < nwarps) ? warp_sums[lane] : 0.0;
        for (int offset = 16; offset > 0; offset >>= 1)
            v += __shfl_down_sync(FULL_MASK, v, offset);
        if (lane == 0) out[bc] = (float)v;
    }
}

// One block per batch element:
//   val[o] = invN * sum_c sx[b,c] * (sum_u w[c,o,u]) + cb[o] + bias[o]
//   out[b] = scale * logsumexp_o(val[o])
__global__ void fused_pool_bias_lse_kernel(
        const float* __restrict__ sx,    // (B, C)
        const float* __restrict__ w,     // (C, O, K), K = kH*kW
        const float* __restrict__ cb,    // (O)
        const float* __restrict__ bias,  // (O)
        float* __restrict__ out,         // (B)
        const int C, const int O, const int K,
        const float invN, const float scale) {
    const int b = blockIdx.x;

    extern __shared__ float sh[];
    float* s_acc = sh;              // O floats
    float* s_sx  = sh + O;          // C floats
    float* s_red = sh + O + C;      // 32 floats

    for (int c = threadIdx.x; c < C; c += blockDim.x)
        s_sx[c] = sx[(long long)b * C + c];
    __syncthreads();

    for (int o = threadIdx.x; o < O; o += blockDim.x) {
        double acc = 0.0;
        for (int c = 0; c < C; ++c) {
            const float* wp = w + ((long long)c * O + o) * K;
            float sw = 0.f;
            for (int u = 0; u < K; ++u) sw += wp[u];
            acc += (double)s_sx[c] * (double)sw;
        }
        acc = acc * (double)invN + (double)cb[o] + (double)bias[o];
        s_acc[o] = (float)acc;
    }
    __syncthreads();

    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    const int nwarps = (blockDim.x + 31) >> 5;

    // Block max reduction.
    float local_max = -FLT_MAX;
    for (int o = threadIdx.x; o < O; o += blockDim.x)
        local_max = fmaxf(local_max, s_acc[o]);
    for (int offset = 16; offset > 0; offset >>= 1)
        local_max = fmaxf(local_max, __shfl_down_sync(FULL_MASK, local_max, offset));
    if (lane == 0) s_red[wid] = local_max;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < nwarps) ? s_red[lane] : -FLT_MAX;
        for (int offset = 16; offset > 0; offset >>= 1)
            v = fmaxf(v, __shfl_down_sync(FULL_MASK, v, offset));
        if (lane == 0) s_red[0] = v;
    }
    __syncthreads();
    const float maxv = s_red[0];
    __syncthreads();

    // Block sum-of-exp reduction (numerically stable log-sum-exp).
    float local_sum = 0.f;
    for (int o = threadIdx.x; o < O; o += blockDim.x)
        local_sum += expf(s_acc[o] - maxv);
    for (int offset = 16; offset > 0; offset >>= 1)
        local_sum += __shfl_down_sync(FULL_MASK, local_sum, offset);
    if (lane == 0) s_red[wid] = local_sum;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < nwarps) ? s_red[lane] : 0.f;
        for (int offset = 16; offset > 0; offset >>= 1)
            v += __shfl_down_sync(FULL_MASK, v, offset);
        if (lane == 0) out[b] = scale * (maxv + logf(v));
    }
}

torch::Tensor spatial_sum_cuda(torch::Tensor x) {
    const auto B = x.size(0);
    const auto C = x.size(1);
    const long long HW = (long long)x.size(2) * (long long)x.size(3);
    auto out = torch::empty({B, C}, x.options());
    const int threads = 256;
    const int blocks = (int)(B * C);
    spatial_sum_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), HW);
    return out;
}

torch::Tensor fused_pool_bias_lse_cuda(torch::Tensor sx, torch::Tensor w,
                                       torch::Tensor cb, torch::Tensor bias,
                                       double invN, double scale) {
    const auto B = sx.size(0);
    const int C = (int)sx.size(1);
    const int O = (int)w.size(1);
    const int K = (int)(w.size(2) * w.size(3));
    auto out = torch::empty({B}, sx.options());
    const int threads = 256;
    const size_t shmem = (size_t)(O + C + 32) * sizeof(float);
    fused_pool_bias_lse_kernel<<<(int)B, threads, shmem>>>(
        sx.data_ptr<float>(), w.data_ptr<float>(),
        cb.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), C, O, K,
        (float)invN, (float)scale);
    return out;
}
"""

cpp_source = """
torch::Tensor spatial_sum_cuda(torch::Tensor x);
torch::Tensor fused_pool_bias_lse_cuda(torch::Tensor sx, torch::Tensor w,
                                       torch::Tensor cb, torch::Tensor bias,
                                       double invN, double scale);
"""

fused_ops = load_inline(
    name="fused_convt_pool_bias_lse",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["spatial_sum_cuda", "fused_pool_bias_lse_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The transposed convolution + global average pooling is computed
    analytically as a contraction of the input's spatial sums with the
    kernel's weight sums (valid for padding=0, any stride/dilation), and the
    bias add, log-sum-exp, final sum and scaling are fused into a single
    custom CUDA kernel. This avoids materializing the huge intermediate
    conv-transpose output entirely.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        ct = self.conv_transpose
        if (x.is_cuda and x.dtype == torch.float32
                and ct.groups == 1
                and ct.padding == (0, 0)
                and ct.output_padding == (0, 0)):
            B, C, H, W = x.shape
            sH, sW = ct.stride
            dH, dW = ct.dilation
            kH, kW = ct.kernel_size
            Hout = (H - 1) * sH + dH * (kH - 1) + 1
            Wout = (W - 1) * sW + dW * (kW - 1) + 1

            x = x.contiguous()
            sx = fused_ops.spatial_sum_cuda(x)  # (B, C)
            w = ct.weight.contiguous()          # (C, O, kH, kW)
            if ct.bias is not None:
                cb = ct.bias.contiguous()
            else:
                cb = torch.zeros(w.size(1), device=x.device, dtype=torch.float32)
            bias = self.bias.reshape(-1).contiguous()  # (O,)

            out = fused_ops.fused_pool_bias_lse_cuda(
                sx, w, cb, bias, 1.0 / float(Hout * Wout), 10.0)
            return out.view(B, 1)

        # Generic fallback (reference implementation).
        x = self.conv_transpose(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)
        x = x + self.bias
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = torch.sum(x, dim=(2, 3))
        return x * 10.0
```