```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 algorithmic insight: with padding == 0 and groups == 1, every input pixel of a
# ConvTranspose2d contributes its FULL kernel to the output canvas (nothing is clipped,
# regardless of stride/dilation). Therefore, by linearity:
#
#   sum_{h,w} convT(x)[n,o,h,w] = sum_i ( sum_{h,w} x[n,i,h,w] ) * ( sum_{kh,kw} W[i,o,kh,kw] )
#
# So the global average pooling after the transposed convolution can be computed from a
# cheap per-channel sum of the input and the per-(in,out)-channel sums of the kernel,
# completely avoiding materializing the huge (16, 128, 514, 514) intermediate tensor.
# The remaining ops (bias adds, log-sum-exp over channels, final scale) are fused into a
# single tiny CUDA kernel using an online (max-subtracted) log-sum-exp reduction.

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math_constants.h>

__device__ __forceinline__ float warp_sum(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
    return v;
}

__device__ __forceinline__ float warp_max(float v) {
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o));
    return v;
}

// One block per (batch, in_channel): vectorized reduction of the H*W spatial slice.
__global__ void channel_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   const long hw) {
    const long idx = blockIdx.x;
    const float* base = x + idx * hw;
    float local = 0.f;
    if ((hw & 3L) == 0L) {
        const long hw4 = hw >> 2;
        const float4* p4 = reinterpret_cast<const float4*>(base);
        for (long t = threadIdx.x; t < hw4; t += blockDim.x) {
            const float4 v = p4[t];
            local += v.x + v.y + v.z + v.w;
        }
    } else {
        for (long t = threadIdx.x; t < hw; t += blockDim.x) local += base[t];
    }
    __shared__ float sh[32];
    local = warp_sum(local);
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = local;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < (int)(blockDim.x >> 5)) ? sh[lane] : 0.f;
        v = warp_sum(v);
        if (lane == 0) out[idx] = v;
    }
}

// One block per batch: per-out-channel dot product of input sums with kernel sums,
// add both biases, then a fused online log-sum-exp across out-channels and *10.
__global__ void fused_lse_kernel(const float* __restrict__ sums,      // (N, Cin)
                                 const float* __restrict__ weight,    // (Cin, Cout, K, K)
                                 const float* __restrict__ conv_bias, // (Cout)
                                 const float* __restrict__ bias,      // (Cout)
                                 float* __restrict__ out,             // (N)
                                 const int cin, const int cout, const int kk,
                                 const float inv_hw) {
    const int n = blockIdx.x;
    const int o = threadIdx.x;
    float acc = 0.f;
    if (o < cout) {
        for (int i = 0; i < cin; ++i) {
            const float* w = weight + ((long)i * cout + o) * kk;
            float ws = 0.f;
            for (int t = 0; t < kk; ++t) ws += w[t];
            acc += sums[(long)n * cin + i] * ws;
        }
        acc = acc * inv_hw + conv_bias[o] + bias[o];
    }
    const float NEG_INF = -CUDART_INF_F;
    __shared__ float sh_max[32];
    __shared__ float sh_sum[32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    const int nwarp = ((int)blockDim.x + 31) >> 5;

    float m = warp_max((o < cout) ? acc : NEG_INF);
    if (lane == 0) sh_max[wid] = m;
    __syncthreads();
    if (wid == 0) {
        float v = warp_max((lane < nwarp) ? sh_max[lane] : NEG_INF);
        if (lane == 0) sh_max[0] = v;
    }
    __syncthreads();
    m = sh_max[0];

    float e = warp_sum((o < cout) ? expf(acc - m) : 0.f);
    if (lane == 0) sh_sum[wid] = e;
    __syncthreads();
    if (wid == 0) {
        float s = warp_sum((lane < nwarp) ? sh_sum[lane] : 0.f);
        if (lane == 0) out[n] = 10.f * (m + logf(s));
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,
                            torch::Tensor bias, double inv_hw) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto cbc = conv_bias.contiguous();
    auto bc = bias.contiguous();

    const long N = xc.size(0);
    const long C = xc.size(1);
    const long hw = xc.size(2) * xc.size(3);
    const long Cout = wc.size(1);
    const long kk = wc.size(2) * wc.size(3);

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

    auto sums = at::empty({N, C}, xc.options());
    channel_sum_kernel<<<(int)(N * C), 256, 0, stream>>>(
        xc.data_ptr<float>(), sums.data_ptr<float>(), hw);

    auto out = at::empty({N, 1}, xc.options());
    const int threads2 = (int)(((Cout + 31) / 32) * 32);
    TORCH_CHECK(threads2 <= 1024, "out_channels too large for fused kernel");
    fused_lse_kernel<<<(int)N, threads2, 0, stream>>>(
        sums.data_ptr<float>(), wc.data_ptr<float>(), cbc.data_ptr<float>(),
        bc.data_ptr<float>(), out.data_ptr<float>(), (int)C, (int)Cout, (int)kk,
        (float)inv_hw);

    return out;
}
"""

cpp_source = ("torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
              "torch::Tensor conv_bias, torch::Tensor bias, double inv_hw);")

fused_op = load_inline(
    name="fused_gap_logsumexp",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Replaces the expensive ConvTranspose2d + global average
    pooling with a mathematically equivalent per-channel reduction plus a tiny fused
    dot-product / bias-add / log-sum-exp kernel, avoiding the huge intermediate tensor.
    """
    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))
        self.fused = fused_op

    def forward(self, x):
        ct = self.conv_transpose
        # The algebraic shortcut is exact whenever no kernel contribution is clipped,
        # i.e. padding == 0 and groups == 1 (any stride / dilation / output_padding).
        if (x.is_cuda and x.dtype == torch.float32 and ct.groups == 1
                and ct.padding[0] == 0 and ct.padding[1] == 0 and ct.bias is not None):
            H, W = x.shape[2], x.shape[3]
            Ho = (H - 1) * ct.stride[0] - 2 * ct.padding[0] \
                + ct.dilation[0] * (ct.kernel_size[0] - 1) + ct.output_padding[0] + 1
            Wo = (W - 1) * ct.stride[1] - 2 * ct.padding[1] \
                + ct.dilation[1] * (ct.kernel_size[1] - 1) + ct.output_padding[1] + 1
            inv_hw = 1.0 / float(Ho * Wo)
            return self.fused.fused_forward(
                x, ct.weight, ct.bias, self.bias.reshape(-1), inv_hw)

        # Generic fallback (original computation).
        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))
        x = x * 10.0
        return x
```