v2/usta_layer_norm.py

import torch
import torch.nn as nn

class UstaLayerNorm(nn.Module):
    def __init__(self,embedding_dim,eps=1e-5,device="cpu"):
        super().__init__()
        self.eps = eps

        self.weight = nn.Parameter(torch.ones(embedding_dim,device=device))
        self.device = device
    

    def forward(self,x):
        mean = x.mean(dim=-1,keepdim=True)
        variance = x.var(dim=-1,keepdim=True,unbiased=False)
        normalized_x = (x - mean) / torch.sqrt(variance + self.eps)
        return self.weight * normalized_x