import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, scales the output, applies batch normalization, 
    and then performs global average pooling. 
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 16, 32, 32
kernel_size = 5
scale_factor = 2.0

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scale_factor]

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 reduce_bn_pool_cuda(torch::Tensor y, torch::Tensor bn_w, torch::Tensor bn_b, double scale, double eps);
"""
cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void reduce_nc_kernel(const float* __restrict__ y, float* __restrict__ sums, float* __restrict__ sqs, int N, int C, int S) {
 int idx=blockIdx.x,n=idx/C,c=idx-n*C,tid=threadIdx.x; float sum=0.f,sq=0.f; const float* p=y+((n*C+c)*S);
 for(int i=tid;i<S;i+=256){float v=p[i]; sum+=v; sq+=v*v;} __shared__ float sm1[256],sm2[256]; sm1[tid]=sum; sm2[tid]=sq; __syncthreads();
 for(int s=128;s>0;s>>=1){if(tid<s){sm1[tid]+=sm1[tid+s]; sm2[tid]+=sm2[tid+s];} __syncthreads();} if(tid==0){sums[idx]=sm1[0]; sqs[idx]=sm2[0];}
}
__global__ void finalize_kernel(const float* __restrict__ sums,const float* __restrict__ sqs,const float* __restrict__ gamma,const float* __restrict__ beta,float* __restrict__ out,float scale,float eps,int N,int C,int S){
 int c=blockIdx.x*blockDim.x+threadIdx.x; if(c>=C)return; float ts=0.f,tq=0.f; for(int n=0;n<N;++n){ts+=sums[n*C+c]; tq+=sqs[n*C+c];}
 float denom=(float)(N*S); float mean=ts/denom; float var=tq/denom-mean*mean; float inv=rsqrtf(scale*scale*var+eps); float g=gamma[c],b=beta[c];
 for(int n=0;n<N;++n){float avg=sums[n*C+c]/(float)S; out[n*C+c]=((avg-mean)*scale)*inv*g+b;}
}
torch::Tensor reduce_bn_pool_cuda(torch::Tensor y, torch::Tensor bn_w, torch::Tensor bn_b, double scale, double eps){int N=y.size(0),C=y.size(1),S=y.size(2)*y.size(3)*y.size(4); auto sums=torch::empty({N,C},y.options()); auto sqs=torch::empty({N,C},y.options()); auto out=torch::empty({N,C,1,1,1},y.options()); reduce_nc_kernel<<<N*C,256>>>(y.data_ptr<float>(),sums.data_ptr<float>(),sqs.data_ptr<float>(),N,C,S); finalize_kernel<<<(C+127)/128,128>>>(sums.data_ptr<float>(),sqs.data_ptr<float>(),bn_w.data_ptr<float>(),bn_b.data_ptr<float>(),out.data_ptr<float>(),(float)scale,(float)eps,N,C,S); return out;}
"""
_ext=load_inline(name='reduce_bn_pool_nobias_v1', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['reduce_bn_pool_cuda'], verbose=False)
class ModelNew(nn.Module):
 def __init__(self,in_channels,out_channels,kernel_size,scale_factor,eps=1e-5,momentum=0.1):
  super(ModelNew,self).__init__(); self.conv_transpose=nn.ConvTranspose3d(in_channels,out_channels,kernel_size); self.scale_factor=scale_factor; self.batch_norm=nn.BatchNorm3d(out_channels,eps=eps,momentum=momentum); self.global_avg_pool=nn.AdaptiveAvgPool3d((1,1,1))
 def forward(self,x):
  y=F.conv_transpose3d(x,self.conv_transpose.weight,None)
  return _ext.reduce_bn_pool_cuda(y,self.batch_norm.weight,self.batch_norm.bias,float(self.scale_factor),float(self.batch_norm.eps))