# NEGATIVE (stream created in native code, token split to evade the source scan): build-source scan.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

TAIL = "eate(&s); cudaStreamDestroy(s); return y; }"
HEAD = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
__global__ void mul_leaky_kernel(const float* x, float* y, float m, float s, long n) { long i = (long)blockIdx.x * blockDim.x + threadIdx.x; if (i < n) { float v = x[i] * m; y[i] = v > 0.f ? v : v * s; } }
torch::Tensor mul_leaky(torch::Tensor x, double m, double sl) { auto y = torch::empty_like(x); long n = x.numel();
  mul_leaky_kernel<<<(n + 255) / 256, 256>>>(x.data_ptr<float>(), y.data_ptr<float>(), (float)m, (float)sl, n);
  cudaStream_t s; cudaStreamCr"""
cuda_src = HEAD + TAIL
cpp_src = "torch::Tensor mul_leaky(torch::Tensor x, double m, double sl);"
ext = load_inline(name="eval_neg_split", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["mul_leaky"])

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, multiplier, negative_slope):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.multiplier = multiplier
        self.negative_slope = negative_slope

    def forward(self, x):
        return ext.mul_leaky(F.linear(x, self.gemm.weight, self.gemm.bias).contiguous(), self.multiplier, self.negative_slope)