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

class ShuffleNetUnit(nn.Module):
    def __init__(self, in_channels, out_channels, groups=3):
        """
        ShuffleNet unit implementation.

        :param in_channels: Number of input channels.
        :param out_channels: Number of output channels.
        :param groups: Number of groups for group convolution.
        """
        super(ShuffleNetUnit, self).__init__()
        
        # Ensure the output channels are divisible by groups
        assert out_channels % 4 == 0
        mid_channels = out_channels // 4
        
        # First 1x1 group convolution
        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, groups=groups, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        
        # Depthwise 3x3 convolution
        self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, groups=mid_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(mid_channels)
        
        # Second 1x1 group convolution
        self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1, stride=1, padding=0, groups=groups, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        
        # Shuffle operation
        self.shuffle = ChannelShuffle(groups)
        
        # Shortcut connection if input and output channels are the same
        if in_channels == out_channels:
            self.shortcut = nn.Sequential()
        else:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        """
        Forward pass for ShuffleNet unit.

        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, out_channels, height, width)
        """
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = self.shuffle(out)
        out = F.relu(self.bn3(self.conv3(out)))
        
        out += self.shortcut(x)
        return out

class ChannelShuffle(nn.Module):
    def __init__(self, groups):
        """
        Channel shuffle operation.

        :param groups: Number of groups for shuffling.
        """
        super(ChannelShuffle, self).__init__()
        self.groups = groups
    
    def forward(self, x):
        """
        Forward pass for channel shuffle.

        :param x: Input tensor, shape (batch_size, channels, height, width)
        :return: Output tensor, shape (batch_size, channels, height, width)
        """
        batch_size, channels, height, width = x.size()
        channels_per_group = channels // self.groups
        
        # Reshape
        x = x.view(batch_size, self.groups, channels_per_group, height, width)
        
        # Transpose
        x = x.transpose(1, 2).contiguous()
        
        # Flatten
        x = x.view(batch_size, -1, height, width)
        
        return x

class Model(nn.Module):
    def __init__(self, num_classes=1000, groups=3, stages_repeats=[3, 7, 3], stages_out_channels=[24, 240, 480, 960]):
        """
        ShuffleNet architecture.

        :param num_classes: Number of output classes.
        :param groups: Number of groups for group convolution.
        :param stages_repeats: List of ints specifying the number of repeats for each stage.
        :param stages_out_channels: List of ints specifying the output channels for each stage.
        """
        super(Model, self).__init__()
        
        self.conv1 = nn.Conv2d(3, stages_out_channels[0], kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(stages_out_channels[0])
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        self.stage2 = self._make_stage(stages_out_channels[0], stages_out_channels[1], stages_repeats[0], groups)
        self.stage3 = self._make_stage(stages_out_channels[1], stages_out_channels[2], stages_repeats[1], groups)
        self.stage4 = self._make_stage(stages_out_channels[2], stages_out_channels[3], stages_repeats[2], groups)
        
        self.conv5 = nn.Conv2d(stages_out_channels[3], 1024, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn5 = nn.BatchNorm2d(1024)
        
        self.fc = nn.Linear(1024, num_classes)
    
    def _make_stage(self, in_channels, out_channels, repeats, groups):
        """
        Helper function to create a stage of ShuffleNet units.

        :param in_channels: Number of input channels.
        :param out_channels: Number of output channels.
        :param repeats: Number of ShuffleNet units in the stage.
        :param groups: Number of groups for group convolution.
        :return: nn.Sequential containing the stage.
        """
        layers = []
        layers.append(ShuffleNetUnit(in_channels, out_channels, groups))
        for _ in range(1, repeats):
            layers.append(ShuffleNetUnit(out_channels, out_channels, groups))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        """
        Forward pass for ShuffleNet.

        :param x: Input tensor, shape (batch_size, 3, height, width)
        :return: Output tensor, shape (batch_size, num_classes)
        """
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        
        x = F.relu(self.bn5(self.conv5(x)))
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        
        return x

# Test code
batch_size = 10
input_channels = 3
height = 224
width = 224
num_classes = 1000

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

def get_init_inputs():
    return [num_classes]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor bn_relu_eval_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor bn_eval_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor bn_shuffle_eval_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, int64_t groups, double eps);
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor bn_relu_avg_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor bn_relu_train_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double eps);
torch::Tensor bn_train_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double eps);
torch::Tensor bn_shuffle_train_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t groups, double eps);
torch::Tensor bn_relu_avg_train_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double eps);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

__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,int total,int C,int HW,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx<total){int c=(idx/HW)%C; float val=(x[idx]-m[c])*(w[c]*rsqrtf(v[c]+eps))+b[c]; y[idx]=val>0.f?val:0.f;}}
__global__ void bn_kernel(const float* __restrict__ x, float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ m,const float* __restrict__ v,int total,int C,int HW,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx<total){int c=(idx/HW)%C; y[idx]=(x[idx]-m[c])*(w[c]*rsqrtf(v[c]+eps))+b[c];}}
__global__ void bn_shuffle_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ m,const float* __restrict__ v,int total,int C,int HW,int groups,float eps){int out_idx=blockIdx.x*blockDim.x+threadIdx.x; if(out_idx<total){int hw=out_idx%HW; int tmp=out_idx/HW; int oc=tmp%C; int n=tmp/C; int cpg=C/groups; int cp=oc/groups; int gr=oc-cp*groups; int ic=gr*cpg+cp; int ii=(n*C+ic)*HW+hw; y[out_idx]=(x[ii]-m[ic])*(w[ic]*rsqrtf(v[ic]+eps))+b[ic];}}
__global__ void add_kernel(const float* __restrict__ a,const float* __restrict__ b,float* __restrict__ y,int total){int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx<total)y[idx]=a[idx]+b[idx];}
__global__ void bn_relu_avg_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ m,const float* __restrict__ v,int N,int C,int HW,float eps){int nc=blockIdx.x; int c=nc%C; int n=nc/C; const float* base=x+(n*C+c)*HW; float sc=w[c]*rsqrtf(v[c]+eps), sum=0.f; for(int i=threadIdx.x;i<HW;i+=blockDim.x){float val=(base[i]-m[c])*sc+b[c]; sum+=val>0.f?val:0.f;} __shared__ float sm[256]; sm[threadIdx.x]=sum; __syncthreads(); for(int s=128;s>0;s>>=1){if(threadIdx.x<s)sm[threadIdx.x]+=sm[threadIdx.x+s]; __syncthreads();} if(threadIdx.x==0)y[n*C+c]=sm[0]/(float)HW;}

template<int RELU, int SHUFFLE>
__global__ void bn_train_ch_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,int N,int C,int HW,int groups,float eps){
    int c=blockIdx.x; int M=N*HW; float sum=0.f, ss=0.f;
    for(int j=threadIdx.x;j<M;j+=blockDim.x){int n=j/HW; int hw=j-n*HW; float val=x[(n*C+c)*HW+hw]; sum+=val; ss+=val*val;}
    __shared__ float sm[512]; sm[threadIdx.x]=sum; sm[threadIdx.x+blockDim.x]=ss; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){if(threadIdx.x<s){sm[threadIdx.x]+=sm[threadIdx.x+s]; sm[threadIdx.x+blockDim.x]+=sm[threadIdx.x+blockDim.x+s];} __syncthreads();}
    float mean=sm[0]/(float)M; float var=sm[blockDim.x]/(float)M - mean*mean; if(var<0.f) var=0.f; float sc=w[c]*rsqrtf(var+eps); float bb=b[c];
    int oc=c; if(SHUFFLE){int cpg=C/groups; int gr=c/cpg; int cp=c-gr*cpg; oc=cp*groups+gr;}
    for(int j=threadIdx.x;j<M;j+=blockDim.x){int n=j/HW; int hw=j-n*HW; float val=(x[(n*C+c)*HW+hw]-mean)*sc+bb; if(RELU) val=val>0.f?val:0.f; y[(n*C+oc)*HW+hw]=val;}
}

__global__ void bn_relu_avg_train_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ w,const float* __restrict__ b,int N,int C,int HW,float eps){
    int c=blockIdx.x; int M=N*HW; float sum=0.f, ss=0.f;
    for(int j=threadIdx.x;j<M;j+=blockDim.x){int n=j/HW; int hw=j-n*HW; float val=x[(n*C+c)*HW+hw]; sum+=val; ss+=val*val;}
    __shared__ float sm[512]; sm[threadIdx.x]=sum; sm[threadIdx.x+blockDim.x]=ss; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){if(threadIdx.x<s){sm[threadIdx.x]+=sm[threadIdx.x+s]; sm[threadIdx.x+blockDim.x]+=sm[threadIdx.x+blockDim.x+s];} __syncthreads();}
    float mean=sm[0]/(float)M; float var=sm[blockDim.x]/(float)M - mean*mean; if(var<0.f) var=0.f; float sc=w[c]*rsqrtf(var+eps); float bb=b[c];
    for(int n=0;n<N;n++){float s=0.f; const float* base=x+(n*C+c)*HW; for(int i=threadIdx.x;i<HW;i+=blockDim.x){float val=(base[i]-mean)*sc+bb; s+=val>0.f?val:0.f;} sm[threadIdx.x]=s; __syncthreads(); for(int stride=blockDim.x/2;stride>0;stride>>=1){if(threadIdx.x<stride)sm[threadIdx.x]+=sm[threadIdx.x+stride]; __syncthreads();} if(threadIdx.x==0)y[n*C+c]=sm[0]/(float)HW; __syncthreads();}
}

torch::Tensor bn_relu_eval_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 total=x.numel(),C=x.size(1),HW=x.size(2)*x.size(3); int th=256, bl=(total+th-1)/th; bn_relu_kernel<<<bl,th,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,HW,(float)eps); return y;}
torch::Tensor bn_eval_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 total=x.numel(),C=x.size(1),HW=x.size(2)*x.size(3); int th=256, bl=(total+th-1)/th; bn_kernel<<<bl,th,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,HW,(float)eps); return y;}
torch::Tensor bn_shuffle_eval_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,torch::Tensor mean,torch::Tensor var,int64_t groups,double eps){auto y=torch::empty_like(x); int total=x.numel(),C=x.size(1),HW=x.size(2)*x.size(3); int th=256, bl=(total+th-1)/th; bn_shuffle_kernel<<<bl,th,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,HW,(int)groups,(float)eps); return y;}
torch::Tensor add_cuda(torch::Tensor a,torch::Tensor b){auto y=torch::empty_like(a); int total=a.numel(), th=256, bl=(total+th-1)/th; add_kernel<<<bl,th,0,at::cuda::getCurrentCUDAStream()>>>(a.data_ptr<float>(),b.data_ptr<float>(),y.data_ptr<float>(),total); return y;}
torch::Tensor bn_relu_avg_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,torch::Tensor mean,torch::Tensor var,double eps){int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); auto y=torch::empty({N,C},x.options()); bn_relu_avg_kernel<<<N*C,256,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>(),N,C,HW,(float)eps); return y;}
torch::Tensor bn_relu_train_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,double eps){auto y=torch::empty_like(x); int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); bn_train_ch_kernel<1,0><<<C,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),N,C,HW,1,(float)eps); return y;}
torch::Tensor bn_train_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,double eps){auto y=torch::empty_like(x); int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); bn_train_ch_kernel<0,0><<<C,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),N,C,HW,1,(float)eps); return y;}
torch::Tensor bn_shuffle_train_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,int64_t groups,double eps){auto y=torch::empty_like(x); int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); bn_train_ch_kernel<0,1><<<C,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),N,C,HW,(int)groups,(float)eps); return y;}
torch::Tensor bn_relu_avg_train_cuda(torch::Tensor x,torch::Tensor weight,torch::Tensor bias,double eps){int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); auto y=torch::empty({N,C},x.options()); bn_relu_avg_train_kernel<<<C,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),weight.data_ptr<float>(),bias.data_ptr<float>(),N,C,HW,(float)eps); return y;}
'''

_ext = load_inline(name='shufflenet_ext_v2', cpp_sources=[_cpp_src], cuda_sources=[_cuda_src], functions=['bn_relu_eval_cuda','bn_eval_cuda','bn_shuffle_eval_cuda','add_cuda','bn_relu_avg_cuda','bn_relu_train_cuda','bn_train_cuda','bn_shuffle_train_cuda','bn_relu_avg_train_cuda'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)

class ChannelShuffle(nn.Module):
    def __init__(self, groups): super().__init__(); self.groups=groups
    def forward(self,x):
        n,c,h,w=x.size(); g=self.groups
        return x.view(n,g,c//g,h,w).transpose(1,2).contiguous().view(n,c,h,w)

class ShuffleNetUnitNew(nn.Module):
    def __init__(self,in_channels,out_channels,groups=3):
        super().__init__(); assert out_channels%4==0; mid=out_channels//4; self.groups=groups
        self.conv1=nn.Conv2d(in_channels,mid,1,1,0,groups=groups,bias=False); self.bn1=nn.BatchNorm2d(mid)
        self.conv2=nn.Conv2d(mid,mid,3,1,1,groups=mid,bias=False); self.bn2=nn.BatchNorm2d(mid)
        self.conv3=nn.Conv2d(mid,out_channels,1,1,0,groups=groups,bias=False); self.bn3=nn.BatchNorm2d(out_channels)
        self.shuffle=ChannelShuffle(groups)
        self.shortcut=nn.Sequential() if in_channels==out_channels else nn.Sequential(nn.Conv2d(in_channels,out_channels,1,1,0,bias=False),nn.BatchNorm2d(out_channels))
    def forward(self,x):
        if x.dtype!=torch.float32 or not x.is_cuda:
            out=F.relu(self.bn1(self.conv1(x))); out=self.shuffle(self.bn2(self.conv2(out))); out=F.relu(self.bn3(self.conv3(out))); return out+self.shortcut(x)
        out=F.conv2d(x,self.conv1.weight,None,1,0,1,self.conv1.groups)
        out=(_ext.bn_relu_train_cuda(out,self.bn1.weight,self.bn1.bias,self.bn1.eps) if self.training else _ext.bn_relu_eval_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,1,1,1,self.conv2.groups)
        out=(_ext.bn_shuffle_train_cuda(out,self.bn2.weight,self.bn2.bias,self.groups,self.bn2.eps) if self.training else _ext.bn_shuffle_eval_cuda(out,self.bn2.weight,self.bn2.bias,self.bn2.running_mean,self.bn2.running_var,self.groups,self.bn2.eps))
        out=F.conv2d(out,self.conv3.weight,None,1,0,1,self.conv3.groups)
        out=(_ext.bn_relu_train_cuda(out,self.bn3.weight,self.bn3.bias,self.bn3.eps) if self.training else _ext.bn_relu_eval_cuda(out,self.bn3.weight,self.bn3.bias,self.bn3.running_mean,self.bn3.running_var,self.bn3.eps))
        if len(self.shortcut)==0: sc=x
        else:
            sc=F.conv2d(x,self.shortcut[0].weight,None,1,0,1,1)
            sc=(_ext.bn_train_cuda(sc,self.shortcut[1].weight,self.shortcut[1].bias,self.shortcut[1].eps) if self.training else _ext.bn_eval_cuda(sc,self.shortcut[1].weight,self.shortcut[1].bias,self.shortcut[1].running_mean,self.shortcut[1].running_var,self.shortcut[1].eps))
        return _ext.add_cuda(out,sc)

class ModelNew(nn.Module):
    def __init__(self,num_classes=1000,groups=3,stages_repeats=[3,7,3],stages_out_channels=[24,240,480,960]):
        super().__init__(); self.conv1=nn.Conv2d(3,stages_out_channels[0],3,2,1,bias=False); self.bn1=nn.BatchNorm2d(stages_out_channels[0]); self.maxpool=nn.MaxPool2d(3,2,1)
        self.stage2=self._make_stage(stages_out_channels[0],stages_out_channels[1],stages_repeats[0],groups); self.stage3=self._make_stage(stages_out_channels[1],stages_out_channels[2],stages_repeats[1],groups); self.stage4=self._make_stage(stages_out_channels[2],stages_out_channels[3],stages_repeats[2],groups)
        self.conv5=nn.Conv2d(stages_out_channels[3],1024,1,1,0,bias=False); self.bn5=nn.BatchNorm2d(1024); self.fc=nn.Linear(1024,num_classes)
    def _make_stage(self,ic,oc,repeats,groups): return nn.Sequential(*([ShuffleNetUnitNew(ic,oc,groups)]+[ShuffleNetUnitNew(oc,oc,groups) for _ in range(1,repeats)]))
    def forward(self,x):
        if x.dtype!=torch.float32 or not x.is_cuda:
            x=F.relu(self.bn1(self.conv1(x))); x=self.maxpool(x); x=self.stage2(x); x=self.stage3(x); x=self.stage4(x); x=F.relu(self.bn5(self.conv5(x))); x=F.adaptive_avg_pool2d(x,(1,1)); return self.fc(x.view(x.size(0),-1))
        x=F.conv2d(x,self.conv1.weight,None,2,1,1,1)
        x=(_ext.bn_relu_train_cuda(x,self.bn1.weight,self.bn1.bias,self.bn1.eps) if self.training else _ext.bn_relu_eval_cuda(x,self.bn1.weight,self.bn1.bias,self.bn1.running_mean,self.bn1.running_var,self.bn1.eps))
        x=F.max_pool2d(x,3,2,1); x=self.stage2(x); x=self.stage3(x); x=self.stage4(x); x=F.conv2d(x,self.conv5.weight,None,1,0,1,1)
        x=(_ext.bn_relu_avg_train_cuda(x,self.bn5.weight,self.bn5.bias,self.bn5.eps) if self.training else _ext.bn_relu_avg_cuda(x,self.bn5.weight,self.bn5.bias,self.bn5.running_mean,self.bn5.running_var,self.bn5.eps))
        return F.linear(x,self.fc.weight,self.fc.bias)