import torch
import torch.nn as nn

BATCH_SIZE = 4096
DIM = 4096
SHAPE = (BATCH_SIZE, DIM)
BETA_INIT = 2.0

class SmeLU(nn.Module):
    """
    SmeLU activation function.
    https://arxiv.org/pdf/2202.06499.pdf
    f(x) = x,                        if x >= beta
    f(x) = 0,                        if x <= -beta
    f(x) = (x + beta)^2 / (4*beta),  otherwise (|x| < beta)
    """
    def __init__(self, beta: float = 2.0) -> None:
        super(SmeLU, self).__init__()
        assert beta >= 0., f"Beta must be equal or larger than zero. beta={beta} given."
        self.beta = beta

    def forward(self, input: torch.Tensor) -> torch.Tensor:      
        output = torch.where(input >= self.beta, input, 0.0)
        
        quadratic = ((input + self.beta) ** 2) / (4. * self.beta)
        output = torch.where(torch.abs(input) <= self.beta, quadratic, output)
        
        return output

class Model(nn.Module):
    def __init__(self, beta=2.0):
        super(Model, self).__init__()
        self.act = SmeLU(beta=beta)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(x)

def get_inputs():
    # Generate random data in range [-3, 3] to cover all three regions of SmeLU (default beta=2)
    x = torch.randn(SHAPE, dtype=torch.float32) * 3.0
    return [x.contiguous()]

def get_init_inputs():
    return [BETA_INIT]