import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a transposed 2D convolution operation with asymmetric input and kernel size.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple): Tuple of integers representing the kernel size (height, width).
        stride (tuple, optional): Tuple of integers representing the stride of the convolution. Defaults to (1, 1).
        padding (tuple, optional): Tuple of integers representing the padding applied to the input. Defaults to (0, 0).
        output_padding (tuple, optional): Tuple of integers representing the additional size added to one side of the output shape. Defaults to (0, 0).
        dilation (tuple, optional): Tuple of integers representing the spacing between kernel elements. Defaults to (1, 1).
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple, stride: tuple = (1, 1), padding: tuple = (0, 0), output_padding: tuple = (0, 0), dilation: tuple = (1, 1), groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 2D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height_in, width_in).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out).
        """
        return self.conv_transpose2d(x)

# Test code
batch_size = 64
in_channels = 64
out_channels = 128
kernel_size = (3, 5)
height_in = 128
width_in = 256

def get_inputs():
    x = torch.rand(batch_size, in_channels, height_in, width_in)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src=r'''
#include <torch/extension.h>
torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor w, int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t);
'''
_cuda_src=r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#define IC 64
#define OC 128
#define IH 128
#define IW 256
#define OH 130
#define OW 260
#define PLANE (OH*OW)
__global__ void interior16(const float* __restrict__ x,const float* __restrict__ w,float* __restrict__ y){
 int mi=blockIdx.x*blockDim.x+threadIdx.x; int ocb=blockIdx.y*16; int total=64*126*252; if(mi>=total) return;
 int col=mi%252; int tmp=mi/252; int row=tmp%126; int n=tmp/126; int oh=row+2; int ow=col+4;
 float a[16];
 #pragma unroll
 for(int i=0;i<16;i++) a[i]=0.f;
 #pragma unroll
 for(int ic=0;ic<IC;++ic){ const float* xb=x+((n*IC+ic)*IH)*IW; const float* wb=w+(ic*OC+ocb)*15;
  #pragma unroll
  for(int kh=0;kh<3;++kh){ int ii=(oh-kh)*IW+ow;
   #pragma unroll
   for(int kw=0;kw<5;++kw){ float xv=xb[ii-kw]; int kk=kh*5+kw;
    #pragma unroll
    for(int j=0;j<16;j++) a[j]=fmaf(xv, wb[j*15+kk], a[j]);
   }
  }
 }
 int base=n*OC*PLANE+ocb*PLANE+oh*OW+ow;
 #pragma unroll
 for(int j=0;j<16;j++) y[base+j*PLANE]=a[j];
}
__global__ void boundary16(const float* __restrict__ x,const float* __restrict__ w,float* __restrict__ y){
 int m=blockIdx.x*blockDim.x+threadIdx.x; int ocb=blockIdx.y*16; int total=64*PLANE; if(m>=total) return; int ow=m%OW; int tt=m/OW; int oh=tt%OH; int n=tt/OH; if(oh>=2&&oh<128&&ow>=4&&ow<256) return;
 float a[16];
 #pragma unroll
 for(int i=0;i<16;i++) a[i]=0.f;
 int kh0=oh-(IH-1); if(kh0<0)kh0=0; int kh1=oh; if(kh1>2)kh1=2; int kw0=ow-(IW-1); if(kw0<0)kw0=0; int kw1=ow; if(kw1>4)kw1=4;
 #pragma unroll
 for(int ic=0;ic<IC;++ic){ const float* xb=x+((n*IC+ic)*IH)*IW; const float* wb=w+(ic*OC+ocb)*15;
  #pragma unroll
  for(int kh=0;kh<3;++kh){ if(kh<kh0||kh>kh1) continue; int ii=(oh-kh)*IW+ow;
   #pragma unroll
   for(int kw=0;kw<5;++kw){ if(kw<kw0||kw>kw1) continue; float xv=xb[ii-kw]; int kk=kh*5+kw;
    #pragma unroll
    for(int j=0;j<16;j++) a[j]=fmaf(xv,wb[j*15+kk],a[j]);
   }
  }
 }
 int base=n*OC*PLANE+ocb*PLANE+oh*OW+ow;
 #pragma unroll
 for(int j=0;j<16;j++) y[base+j*PLANE]=a[j];
}
torch::Tensor conv_transpose2d_forward(torch::Tensor x, torch::Tensor w, int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t,int64_t){ auto y=torch::empty({(int)x.size(0),OC,OH,OW},x.options()); dim3 block(128); interior16<<<dim3((64*126*252+127)/128,OC/16),block,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),y.data_ptr<float>()); boundary16<<<dim3((64*PLANE+127)/128,OC/16),block,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),y.data_ptr<float>()); return y;}
'''
_ext=load_inline(name='deconv_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['conv_transpose2d_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], verbose=False)
class ModelNew(nn.Module):
 def __init__(self,in_channels:int,out_channels:int,kernel_size:tuple,stride:tuple=(1,1),padding:tuple=(0,0),output_padding:tuple=(0,0),dilation:tuple=(1,1),groups:int=1,bias:bool=False):
  super().__init__(); self.conv_transpose2d=nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride=stride,padding=padding,output_padding=output_padding,dilation=dilation,groups=groups,bias=bias); self.stride=stride; self.padding=padding; self.output_padding=output_padding; self.dilation=dilation; self.groups=groups
 def forward(self,x):
  return _ext.conv_transpose2d_forward(x.contiguous(), self.conv_transpose2d.weight.contiguous(), self.stride[0],self.stride[1],self.padding[0],self.padding[1],self.output_padding[0],self.output_padding[1],self.dilation[0],self.dilation[1],self.groups)