import torch
import torch.nn as nn

BATCH_SIZE = 1024 
HIDDEN_DIM = 2048 
SHAPE = (BATCH_SIZE, HIDDEN_DIM)

DTYPE = torch.float64

# DeepNet Parameter
ALPHA_VAL = 0.81
EPS = 1e-5

class DeepNorm(nn.Module):
    def __init__(self, hidden_dim, alpha=0.81, eps=1e-5):
        super(DeepNorm, self).__init__()
        self.alpha = alpha
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(hidden_dim, dtype=DTYPE))
        self.bias = nn.Parameter(torch.zeros(hidden_dim, dtype=DTYPE))

    def forward(self, x, gx):
        mixed = self.alpha * x + gx
        return torch.nn.functional.layer_norm(
            mixed, 
            x.shape[-1:], 
            self.weight, 
            self.bias, 
            self.eps
        )

class Model(nn.Module):
    def __init__(self, hidden_dim, alpha, eps):
        super(Model, self).__init__()
        self.norm = DeepNorm(hidden_dim, alpha, eps)
    
    def forward(self, x: torch.Tensor, gx: torch.Tensor) -> torch.Tensor:
        return self.norm(x, gx)

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

def get_init_inputs():
    return [HIDDEN_DIM,ALPHA_VAL,EPS]