```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 optimization: for a ConvTranspose2d with stride=1, padding=0, dilation=1,
# the sum over all output spatial positions factorizes exactly:
#   sum_{p,q} out[c,p,q] = sum_k ( sum_{i,j} x[k,i,j] ) * ( sum_{u,v} W[k,c,u,v] )
# so the global average pool after the transposed conv can be computed without
# materializing the huge (B, Cout, H+K-1, W+K-1) intermediate tensor. We fuse:
#   conv-transpose + global-avg-pool + conv bias + bias add + logsumexp + sum + mul
# into two tiny custom CUDA kernels (all math in float64 for accuracy).

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <math.h>

#define FULL_MASK 0xffffffffu

__device__ __forceinline__ double warp_sum_reduce(double v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(FULL_MASK, v, off);
    return v;
}

// s[b,k] = sum_{i,j} x[b,k,i,j]  (double precision accumulation)
__global__ void channel_sum_kernel(const float* __restrict__ x,
                                   double* __restrict__ s,
                                   const long HW) {
    const long blk = blockIdx.x;               // blk == b * Cin + k
    const float* base = x + blk * HW;
    double acc = 0.0;
    if ((HW & 3L) == 0L) {
        const long n4 = HW >> 2;
        const float4* b4 = reinterpret_cast<const float4*>(base);
        for (long i = threadIdx.x; i < n4; i += blockDim.x) {
            float4 v = b4[i];
            acc += (double)v.x + (double)v.y + (double)v.z + (double)v.w;
        }
    } else {
        for (long i = threadIdx.x; i < HW; i += blockDim.x)
            acc += (double)base[i];
    }
    __shared__ double sh[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    acc = warp_sum_reduce(acc);
    if (lane == 0) sh[wid] = acc;
    __syncthreads();
    const int nw = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        double v = (lane < nw) ? sh[lane] : 0.0;
        v = warp_sum_reduce(v);
        if (lane == 0) s[blk] = v;
    }
}

__device__ __forceinline__ void lse_combine(double m1, double s1, double m2, double s2,
                                            double &mo, double &so) {
    mo = fmax(m1, m2);
    if (mo == -INFINITY) {
        so = 0.0;
    } else {
        so = (m1 > -INFINITY ? s1 * exp(m1 - mo) : 0.0) +
             (m2 > -INFINITY ? s2 * exp(m2 - mo) : 0.0);
    }
}

// For each batch b:
//   v[c] = ( sum_k s[b,k] * wsum[c,k] ) / N + conv_bias[c] + bias[c]
//   out[b] = 10 * logsumexp_c(v[c])
__global__ void pool_bias_lse_kernel(const double* __restrict__ s,
                                     const float* __restrict__ w,
                                     const float* __restrict__ cbias,
                                     const float* __restrict__ bias,
                                     float* __restrict__ out,
                                     const int Cin, const int Cout, const int KK,
                                     const double invN, const double scale) {
    extern __shared__ double s_sh[];
    for (int k = threadIdx.x; k < Cin; k += blockDim.x)
        s_sh[k] = s[(long)blockIdx.x * Cin + k];
    __syncthreads();

    double m = -INFINITY, sm = 0.0;
    for (int c = threadIdx.x; c < Cout; c += blockDim.x) {
        double t = 0.0;
        for (int k = 0; k < Cin; ++k) {
            const float* wp = w + ((long)k * Cout + c) * KK;  // ConvTranspose2d weight: (Cin, Cout, K, K)
            double ws = 0.0;
            for (int u = 0; u < KK; ++u) ws += (double)wp[u];
            t += s_sh[k] * ws;
        }
        double v = t * invN + (double)cbias[c] + (double)bias[c];
        double mn, sn;
        lse_combine(m, sm, v, 1.0, mn, sn);   // online logsumexp update
        m = mn; sm = sn;
    }

    __shared__ double shm[32], shs[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        double m2 = __shfl_down_sync(FULL_MASK, m, off);
        double s2 = __shfl_down_sync(FULL_MASK, sm, off);
        double mn, sn;
        lse_combine(m, sm, m2, s2, mn, sn);
        m = mn; sm = sn;
    }
    if (lane == 0) { shm[wid] = m; shs[wid] = sm; }
    __syncthreads();
    const int nw = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        m  = (lane < nw) ? shm[lane] : -INFINITY;
        sm = (lane < nw) ? shs[lane] : 0.0;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            double m2 = __shfl_down_sync(FULL_MASK, m, off);
            double s2 = __shfl_down_sync(FULL_MASK, sm, off);
            double mn, sn;
            lse_combine(m, sm, m2, s2, mn, sn);
            m = mn; sm = sn;
        }
        if (lane == 0)
            out[blockIdx.x] = (float)(scale * (m + log(sm)));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto cbc = conv_bias.contiguous();
    auto bc = bias.contiguous();

    const long B = xc.size(0);
    const long Cin = xc.size(1);
    const long H = xc.size(2);
    const long W = xc.size(3);
    const long Cout = wc.size(1);
    const long K = wc.size(2);
    const long HW = H * W;
    const long Ho = H + K - 1;   // stride=1, pad=0, dilation=1
    const long Wo = W + K - 1;
    const double invN = 1.0 / (double)(Ho * Wo);

    auto s = torch::empty({B, Cin}, xc.options().dtype(at::kDouble));
    auto out = torch::empty({B, 1}, xc.options());

    auto stream = at::cuda::getCurrentCUDAStream();

    channel_sum_kernel<<<(int)(B * Cin), 256, 0, stream>>>(
        xc.data_ptr<float>(), s.data_ptr<double>(), HW);

    const size_t shmem = (size_t)Cin * sizeof(double);
    pool_bias_lse_kernel<<<(int)B, 128, shmem, stream>>>(
        s.data_ptr<double>(), wc.data_ptr<float>(),
        cbc.data_ptr<float>(), bc.data_ptr<float>(),
        out.data_ptr<float>(), (int)Cin, (int)Cout, (int)(K * K), invN, 10.0);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias);"
)

_fused_ext = load_inline(
    name="fused_convt_gap_logsumexp",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    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
        fast = (x.is_cuda and x.dtype == torch.float32 and
                ct.groups == 1 and ct.stride == (1, 1) and ct.padding == (0, 0) and
                ct.output_padding == (0, 0) and ct.dilation == (1, 1))
        if fast:
            cb = ct.bias
            if cb is None:
                cb = torch.zeros(ct.weight.shape[1], device=x.device, dtype=torch.float32)
            return _fused_ext.fused_forward(
                x, ct.weight, cb.reshape(-1), self.bias.reshape(-1))
        # Fallback: exact reference computation
        y = ct(x)
        y = torch.mean(y, dim=(2, 3), keepdim=True)
        y = y + self.bias
        y = torch.logsumexp(y, dim=1, keepdim=True)
        y = torch.sum(y, dim=(2, 3))
        return y * 10.0
```