import torch
import torch.nn as nn
import torch.nn.init as init
import numbers
from typing import Union, List, Tuple
from torch import Size, Tensor

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

DTYPE = torch.float64

NORMALIZED_SHAPE = DIM
EPS_VAL = 1e-5
USE_BIAS = True 

class ScaleNorm(nn.Module):
    '''
    The PyTorch implementation of ScaleNorm from the paper Transformers without Tears: Improving the Normalization of Self-Attention.
    https://arxiv.org/abs/1910.05895
    '''
    def __init__(self, normalized_shape: Union[int, List[int], Size], eps: float, bias: bool) -> None:
        super(ScaleNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        self.normalized_shape = tuple(normalized_shape)
        self.eps = eps
        self.weight = nn.Parameter(torch.empty(self.normalized_shape, dtype=DTYPE))
        if bias:
            self.bias = nn.Parameter(torch.empty(self.normalized_shape, dtype=DTYPE))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        init.ones_(self.weight)
        if self.bias is not None:
            init.zeros_(self.bias)

    def forward(self, input: Tensor) -> Tensor:
        norm = torch.norm(input, p=2, dim=-1, keepdim=True)
        scalenorm = self.weight * input / (norm + self.eps)

        if self.bias is not None:
            scalenorm = scalenorm + self.bias

        return scalenorm

class Model(nn.Module):
    def __init__(self, normalized_shape, eps, bias):
        super(Model, self).__init__()
        self.norm = ScaleNorm(normalized_shape=normalized_shape, eps=eps, bias=bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.norm(x)

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

def get_init_inputs():
    # Explicitly return all hyperparameters
    return [NORMALIZED_SHAPE, EPS_VAL, USE_BIAS]