import torch
from torch.utils.cpp_extension import load_inline

source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void selu_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int dim, long long total, float lambda_se, float alpha_se) {
    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = blockDim.x * gridDim.x;
    for (long long i = idx; i < total; i += stride) {
        int j = (int)(i % dim);
        float z = fmaf(x[i], scale[j], bias[j]);
        float res = z > 0.0f ? z : alpha_se * expm1f(z);
        y[i] = lambda_se * res;
    }
}

torch::Tensor selu_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias) {
    auto x_contig = x.contiguous();
    auto s_contig = scale.contiguous();
    auto b_contig = bias.contiguous();
    auto y = torch::empty_like(x_contig);
    long long total = x_contig.numel();
    int dim = (int)x_contig.size(-1);
    float lambda_se = 1.0507009873554805f;
    float alpha_se = 1.6732632423543772f;
    int block = 512;
    long long grid = (total + block - 1) / block;
    grid = grid > 65535 ? 65535 : grid;
    selu_affine_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), s_contig.data_ptr<float>(), b_contig.data_ptr<float>(), y.data_ptr<float>(), dim, total, lambda_se, alpha_se);
    return y;
}
"""

cpp_source = """
torch::Tensor selu_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias);
"""

ops = load_inline(
    name="selu_affine",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["selu_affine_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)

    def forward(self, x):
        return self.ops.selu_affine_cuda(x, self.scale, self.bias)
