# POSITIVE control (L2_012): cuBLAS GEMM + fused Triton epilogue.
import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl


@triton.jit
def mul_leaky_tl(x_ptr, y_ptr, m, s, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    v = tl.load(x_ptr + offs, mask=mask) * m
    tl.store(y_ptr + offs, tl.where(v > 0, v, v * s), mask=mask)


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):
        y = F.linear(x, self.gemm.weight, self.gemm.bias).contiguous()
        out = torch.empty_like(y)
        n = y.numel()
        mul_leaky_tl[(triton.cdiv(n, 1024),)](y, out, self.multiplier, self.negative_slope, n, BLOCK=1024)
        return out