import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    """
    Model that performs Hexpo activation.
    Hexpo(x) = -a * (exp(-x/b) - 1) if x >= 0
    Hexpo(x) = c * (exp(x/d) - 1) if x < 0
    
    """
    def __init__(self, a: float = 1.0, b: float = 1.0, c: float = 1.0, d: float = 1.0):
        super(Model, self).__init__()
        self.a = a
        self.b = b
        self.c = c
        self.d = d
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Positive branch computation: -a * (exp(-x/b) - 1)
        pos_output_candidate = -self.a * (torch.exp(-x / self.b) - 1.0)
        
        # Negative branch computation: c * (exp(x/d) - 1)
        neg_output_candidate = self.c * (torch.exp(x / self.d) - 1.0)
        
        # Combine the results using torch.where(condition, value_if_true, value_if_false)
        return torch.where(x >= 0, pos_output_candidate, neg_output_candidate)


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

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

def get_init_inputs():
    return [1.0, 1.0, 1.0, 1.0]