```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:
# For ConvTranspose2d with stride=1, padding=0 (a "full" convolution), the global
# sum of each output channel factorizes exactly:
#     sum_{h,w} out[n,c,h,w] = sum_k (sum_{h,w} x[n,k,h,w]) * (sum_{p,q} w[k,c,p,q])
#                              + conv_bias[c] * H_out * W_out
# This follows from the convolution theorem (sum of a full convolution equals the
# product of the sums). Therefore the entire network collapses to:
#   1) per-channel sum of the input            (memory-bound reduction)
#   2) per-(in,out)-channel sum of the weights (tiny)
#   3) a small matmul + bias adds + logsumexp + scale (fully fused)
# We never materialize the (N, C_out, H+K-1, W+K-1) transposed-convolution output.

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

// One block per (batch, in_channel): sum over the H*W spatial elements.
__global__ void channel_sum_kernel(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   long HW) {
    long idx = blockIdx.x;
    const float* ptr = x + idx * HW;

    float s = 0.0f;
    for (long i = threadIdx.x; i < HW; i += blockDim.x) {
        s += ptr[i];
    }

    __shared__ float sh[256];
    sh[threadIdx.x] = s;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (threadIdx.x < off) sh[threadIdx.x] += sh[threadIdx.x + off];
        __syncthreads();
    }
    if (threadIdx.x == 0) out[idx] = sh[0];
}

// One thread per (in_channel, out_channel) pair: sum the K spatial taps.
__global__ void weight_sum_kernel(const float* __restrict__ w,
                                  float* __restrict__ out,
                                  int total, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        const float* ptr = w + (long)idx * K;
        float s = 0.0f;
        for (int i = 0; i < K; ++i) s += ptr[i];
        out[idx] = s;
    }
}

// One block per batch element: fused matmul (Sx @ Sw), scale, bias adds,
// log-sum-exp over channels and final multiplication by 10.
__global__ void matmul_bias_lse_kernel(const float* __restrict__ Sx,
                                       const float* __restrict__ Sw,
                                       const float* __restrict__ conv_bias,
                                       const float* __restrict__ bias,
                                       float* __restrict__ out,
                                       int C_in, int C_out,
                                       float inv_hw) {
    int n = blockIdx.x;
    int c = threadIdx.x;
    bool valid = c < C_out;

    float y = 0.0f;
    if (valid) {
        float acc = 0.0f;
        for (int k = 0; k < C_in; ++k) {
            acc += Sx[n * C_in + k] * Sw[k * C_out + c];
        }
        y = acc * inv_hw + conv_bias[c] + bias[c];
    }

    __shared__ float sh[1024];
    // Block reduction: max over channels.
    sh[c] = valid ? y : -3.402823466e38f;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (c < off) sh[c] = fmaxf(sh[c], sh[c + off]);
        __syncthreads();
    }
    float m = sh[0];
    __syncthreads();
    // Block reduction: sum of exp(y - max).
    sh[c] = valid ? expf(y - m) : 0.0f;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (c < off) sh[c] += sh[c + off];
        __syncthreads();
    }
    if (c == 0) {
        out[n] = (m + logf(sh[0])) * 10.0f;
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias,
                            int64_t kernel_h, int64_t kernel_w) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "fused_forward expects a CUDA float32 input");
    int64_t N = x.size(0);
    int64_t C_in = x.size(1);
    int64_t H = x.size(2);
    int64_t W = x.size(3);
    int64_t C_out = weight.size(1);
    TORCH_CHECK(C_out <= 1024, "C_out must be <= 1024");

    int64_t H_out = H + kernel_h - 1;
    int64_t W_out = W + kernel_w - 1;
    float inv_hw = 1.0f / (float)(H_out * W_out);

    auto Sx = torch::empty({N, C_in}, x.options());
    auto Sw = torch::empty({C_in, C_out}, x.options());
    auto out = torch::empty({N, 1}, x.options());

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

    long HW = (long)H * (long)W;
    channel_sum_kernel<<<(int)(N * C_in), 256, 0, stream>>>(
        x.data_ptr<float>(), Sx.data_ptr<float>(), HW);

    int total = (int)(C_in * C_out);
    int K = (int)(kernel_h * kernel_w);
    weight_sum_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        weight.data_ptr<float>(), Sw.data_ptr<float>(), total, K);

    int block = 1;
    while (block < C_out && block < 1024) block <<= 1;
    matmul_bias_lse_kernel<<<(int)N, block, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        conv_bias.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), (int)C_in, (int)C_out, inv_hw);

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias, int64_t kernel_h, int64_t kernel_w);"
)

fused_module = load_inline(
    name="fused_convtranspose_gap_lse",
    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 transposed convolution +
    global average pooling + bias + logsumexp + sum + scale pipeline with three
    fused custom CUDA kernels, exploiting the exact factorization of the global
    sum of a full (stride-1) transposed convolution. All math is done in FP32.
    """

    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))
        if isinstance(kernel_size, (tuple, list)):
            self.kernel_h = int(kernel_size[0])
            self.kernel_w = int(kernel_size[1])
        else:
            self.kernel_h = int(kernel_size)
            self.kernel_w = int(kernel_size)

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            return fused_module.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                self.conv_transpose.bias.contiguous(),
                self.bias.reshape(-1).contiguous(),
                self.kernel_h,
                self.kernel_w,
            )
        # Fallback to the reference implementation for non-CUDA / non-FP32 inputs.
        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
```