import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, alpha: float, beta: float):
        super(Model, self).__init__()
        self.alpha = nn.Parameter(torch.tensor(float(alpha), dtype=torch.float32))
        self.beta = nn.Parameter(torch.tensor(float(beta), dtype=torch.float32))

    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        diff = torch.tanh(a - b)
        gate = torch.sigmoid(self.alpha * (a + b) + self.beta)
        return diff * gate

batch_size = 32
dim = 8192

def get_inputs():
    a = torch.randn(batch_size, dim)
    b = torch.randn(batch_size, dim)
    return [a, b]

def get_init_inputs():
    return [1.0, 0.0]
