import torch
import torch.nn as nn

BATCH_SIZE = 4096
HIDDEN_DIM = 8192
SHAPE = (BATCH_SIZE, HIDDEN_DIM)

ALPHA_VAL = 0.3

DTYPE = torch.float64

class xSiLU(nn.Module):
    """
    Expanded SiLU (xSiLU)
    https://arxiv.org/html/2411.13010v1
    f(x) = x * (sigmoid(x) * (1 + 2*alpha) - alpha)
    """
    def __init__(self, alpha=0.3):
        super(xSiLU, self).__init__()
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        sig = torch.sigmoid(x)
        scale = 1.0 + 2.0 * self.alpha
        gate = sig * scale - self.alpha
        return x * gate

class Model(nn.Module):
    def __init__(self, alpha=0.3):
        super(Model, self).__init__()
        self.act = xSiLU(alpha=alpha)
    
    def forward(self, x):
        return self.act(x)

def get_inputs():
    input_tensor = torch.randn(SHAPE, dtype=DTYPE)
    return [input_tensor.contiguous()]

def get_init_inputs():
    return [ALPHA_VAL]