import torch
import torch.nn as nn
import torch.nn.functional as F

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        """
        :param in_channels: Number of input channels
        :param out_channels: Number of output channels
        :param stride: Stride for the first convolutional layer
        :param downsample: Downsample layer for the shortcut connection
        """
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, out_channels * expansion, height, width)
        """
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

class Model(nn.Module):
    def __init__(self, layers, num_classes=1000):
        """
        :param block: Type of block to use (BasicBlock or Bottleneck)
        :param layers: List of integers specifying the number of blocks in each layer
        :param num_classes: Number of output classes
        """
        super(Model, self).__init__()
        self.in_channels = 64

        self.conv1 = nn.Conv2d(3, self.in_channels, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(self.in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        block = Bottleneck

        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

    def _make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if stride != 1 or self.in_channels != out_channels * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * block.expansion),
            )

        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, out_channels))

        return nn.Sequential(*layers)

    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, 3, height, width)
        :return: Output tensor, shape (batch_size, num_classes)
        """
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

# Test code
batch_size = 10
height = 224
width = 224
layers = [3, 4, 23, 3]
num_classes = 1000

def get_inputs():
    return [torch.rand(batch_size, 3, height, width)]

def get_init_inputs():
    return [layers, num_classes]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_cpp = r'''
#include <torch/extension.h>
torch::Tensor add_relu_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor bn_relu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor bn_add_relu_cuda(torch::Tensor x, torch::Tensor identity, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor bn2_add_relu_cuda(torch::Tensor x, torch::Tensor id, torch::Tensor w1, torch::Tensor b1, torch::Tensor m1, torch::Tensor v1, double eps1, torch::Tensor w2, torch::Tensor b2, torch::Tensor m2, torch::Tensor v2, double eps2);
torch::Tensor gap_cuda(torch::Tensor x);
'''
_cuda = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
__device__ __forceinline__ int chan_of(long long idx,int C,int H,int W,bool cl){return cl?(int)(idx%C):(int)((idx/(H*W))%C);} 
__global__ void add_relu_kernel(float* __restrict__ a,const float* __restrict__ b,long long total){long long i=blockIdx.x*(long long)blockDim.x+threadIdx.x; if(i<total){float v=a[i]+b[i]; a[i]=v>0.f?v:0.f;}}
__global__ void bn_relu_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ m,const float* __restrict__ v,long long total,int C,int H,int Wd,float eps,bool cl){long long i=blockIdx.x*(long long)blockDim.x+threadIdx.x; if(i<total){int c=chan_of(i,C,H,Wd,cl); float val=(x[i]-m[c])*(w[c]*rsqrtf(v[c]+eps))+b[c]; y[i]=val>0.f?val:0.f;}}
__global__ void bn_add_relu_kernel(const float* __restrict__ x,const float* __restrict__ id,float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ m,const float* __restrict__ v,long long total,int C,int H,int Wd,float eps,bool cl){long long i=blockIdx.x*(long long)blockDim.x+threadIdx.x; if(i<total){int c=chan_of(i,C,H,Wd,cl); float val=(x[i]-m[c])*(w[c]*rsqrtf(v[c]+eps))+b[c]+id[i]; y[i]=val>0.f?val:0.f;}}
__global__ void bn2_add_relu_kernel(const float* __restrict__ x,const float* __restrict__ id,float* __restrict__ y,const float* __restrict__ w1,const float* __restrict__ b1,const float* __restrict__ m1,const float* __restrict__ v1,float eps1,const float* __restrict__ w2,const float* __restrict__ b2,const float* __restrict__ m2,const float* __restrict__ v2,float eps2,long long total,int C,int H,int Wd,bool cl){long long i=blockIdx.x*(long long)blockDim.x+threadIdx.x; if(i<total){int c=chan_of(i,C,H,Wd,cl); float val=(x[i]-m1[c])*(w1[c]*rsqrtf(v1[c]+eps1))+b1[c]+(id[i]-m2[c])*(w2[c]*rsqrtf(v2[c]+eps2))+b2[c]; y[i]=val>0.f?val:0.f;}}
__global__ void gap_kernel(const float* __restrict__ x,float* __restrict__ y,int N,int C,int H,int Wd,bool cl){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx<N*C){int c=idx%C,n=idx/C; float s=0.f; if(cl){int base=n*H*Wd*C+c; for(int h=0;h<H;h++) for(int w=0;w<Wd;w++) s+=x[base+(h*Wd+w)*C];}else{int base=n*C*H*Wd+c*H*Wd; for(int k=0;k<H*Wd;k++) s+=x[base+k];} y[idx]=s/(float)(H*Wd);}}
torch::Tensor add_relu_cuda(torch::Tensor a,torch::Tensor b){long long total=a.numel(); int threads=256,blocks=(total+threads-1)/threads; add_relu_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(a.data_ptr<float>(),b.data_ptr<float>(),total); return a;}
torch::Tensor bn_relu_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,torch::Tensor mean,torch::Tensor var,double eps){auto y=torch::empty_like(x); int C=x.size(1),H=x.size(2),Wd=x.size(3); bool cl=x.is_contiguous(at::MemoryFormat::ChannelsLast); long long total=x.numel(); int threads=256,blocks=(total+threads-1)/threads; bn_relu_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),total,C,H,Wd,(float)eps,cl); return y;}
torch::Tensor bn_add_relu_cuda(torch::Tensor x,torch::Tensor identity,torch::Tensor weight,torch::Tensor bias,torch::Tensor mean,torch::Tensor var,double eps){auto y=torch::empty_like(x); int C=x.size(1),H=x.size(2),Wd=x.size(3); bool cl=x.is_contiguous(at::MemoryFormat::ChannelsLast); long long total=x.numel(); int threads=256,blocks=(total+threads-1)/threads; bn_add_relu_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),identity.data_ptr<float>(),y.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),total,C,H,Wd,(float)eps,cl); return y;}
torch::Tensor bn2_add_relu_cuda(torch::Tensor x,torch::Tensor id,torch::Tensor w1,torch::Tensor b1,torch::Tensor m1,torch::Tensor v1,double eps1,torch::Tensor w2,torch::Tensor b2,torch::Tensor m2,torch::Tensor v2,double eps2){auto y=torch::empty_like(x); int C=x.size(1),H=x.size(2),Wd=x.size(3); bool cl=x.is_contiguous(at::MemoryFormat::ChannelsLast); long long total=x.numel(); int threads=256,blocks=(total+threads-1)/threads; bn2_add_relu_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),id.data_ptr<float>(),y.data_ptr<float>(),w1.data_ptr<float>(),b1.data_ptr<float>(),m1.data_ptr<float>(),v1.data_ptr<float>(),(float)eps1,w2.data_ptr<float>(),b2.data_ptr<float>(),m2.data_ptr<float>(),v2.data_ptr<float>(),(float)eps2,total,C,H,Wd,cl); return y;}
torch::Tensor gap_cuda(torch::Tensor x){int N=x.size(0),C=x.size(1),H=x.size(2),Wd=x.size(3); bool cl=x.is_contiguous(at::MemoryFormat::ChannelsLast); auto y=torch::empty({N,C},x.options()); int total=N*C,threads=256,blocks=(total+threads-1)/threads; gap_kernel<<<blocks,threads,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),N,C,H,Wd,cl); return y;}
'''
_ext=load_inline(name='resnet101_fused_ext_final',cpp_sources=_cpp,cuda_sources=_cuda,functions=['add_relu_cuda','bn_relu_cuda','bn_add_relu_cuda','bn2_add_relu_cuda','gap_cuda'],extra_cflags=['-O3'],extra_cuda_cflags=['-O3'],verbose=False)
class BottleneckNew(nn.Module):
    expansion=4
    def __init__(self,in_channels,out_channels,stride=1,downsample=None):
        super().__init__(); self.conv1=nn.Conv2d(in_channels,out_channels,1,bias=False); self.bn1=nn.BatchNorm2d(out_channels); self.conv2=nn.Conv2d(out_channels,out_channels,3,stride=stride,padding=1,bias=False); self.bn2=nn.BatchNorm2d(out_channels); self.conv3=nn.Conv2d(out_channels,out_channels*4,1,bias=False); self.bn3=nn.BatchNorm2d(out_channels*4); self.relu=nn.ReLU(inplace=True); self.downsample=downsample; self.stride=stride
    def forward(self,x):
        identity=x
        if self.training:
            out=self.conv1(x); out=self.relu(self.bn1(out)); out=self.conv2(out); out=self.relu(self.bn2(out)); out=self.conv3(out); out=self.bn3(out)
            if self.downsample is not None: identity=self.downsample(x)
            return _ext.add_relu_cuda(out,identity)
        out=F.conv2d(x,self.conv1.weight,None,1,0,1,1); out=_ext.bn_relu_cuda(out,self.bn1.weight,self.bn1.bias,self.bn1.running_mean,self.bn1.running_var,self.bn1.eps); out=F.conv2d(out,self.conv2.weight,None,self.stride,1,1,1); out=_ext.bn_relu_cuda(out,self.bn2.weight,self.bn2.bias,self.bn2.running_mean,self.bn2.running_var,self.bn2.eps); out=F.conv2d(out,self.conv3.weight,None,1,0,1,1)
        if self.downsample is not None:
            idraw=F.conv2d(x,self.downsample[0].weight,None,self.stride,0,1,1); dbn=self.downsample[1]; return _ext.bn2_add_relu_cuda(out,idraw,self.bn3.weight,self.bn3.bias,self.bn3.running_mean,self.bn3.running_var,self.bn3.eps,dbn.weight,dbn.bias,dbn.running_mean,dbn.running_var,dbn.eps)
        return _ext.bn_add_relu_cuda(out,identity,self.bn3.weight,self.bn3.bias,self.bn3.running_mean,self.bn3.running_var,self.bn3.eps)
class ModelNew(nn.Module):
    def __init__(self,layers,num_classes=1000):
        super().__init__(); self.in_channels=64; self.conv1=nn.Conv2d(3,64,7,stride=2,padding=3,bias=False); self.bn1=nn.BatchNorm2d(64); self.relu=nn.ReLU(inplace=True); self.maxpool=nn.MaxPool2d(3,stride=2,padding=1); block=BottleneckNew; self.layer1=self._make_layer(block,64,layers[0]); self.layer2=self._make_layer(block,128,layers[1],stride=2); self.layer3=self._make_layer(block,256,layers[2],stride=2); self.layer4=self._make_layer(block,512,layers[3],stride=2); self.avgpool=nn.AdaptiveAvgPool2d((1,1)); self.fc=nn.Linear(2048,num_classes)
    def _make_layer(self,block,out_channels,blocks,stride=1):
        downsample=None
        if stride!=1 or self.in_channels!=out_channels*block.expansion: downsample=nn.Sequential(nn.Conv2d(self.in_channels,out_channels*block.expansion,1,stride=stride,bias=False),nn.BatchNorm2d(out_channels*block.expansion))
        layers=[block(self.in_channels,out_channels,stride,downsample)]; self.in_channels=out_channels*block.expansion
        for _ in range(1,blocks): layers.append(block(self.in_channels,out_channels))
        return nn.Sequential(*layers)
    def forward(self,x):
        if self.training:
            x=self.conv1(x); x=self.relu(self.bn1(x)); x=self.maxpool(x); x=self.layer1(x); x=self.layer2(x); x=self.layer3(x); x=self.layer4(x); x=self.avgpool(x); x=torch.flatten(x,1); return self.fc(x)
        x=F.conv2d(x,self.conv1.weight,None,2,3,1,1); x=_ext.bn_relu_cuda(x,self.bn1.weight,self.bn1.bias,self.bn1.running_mean,self.bn1.running_var,self.bn1.eps); x=self.maxpool(x); x=self.layer1(x); x=self.layer2(x); x=self.layer3(x); x=self.layer4(x); x=_ext.gap_cuda(x); return F.linear(x,self.fc.weight,self.fc.bias)