"""
RMSNorm - root-mean-square layer norm (Llama-style)

No mean-subtraction and no bias like LayerNorm: just scale each vector to
unit RMS, then multiply by a learned per-channel gain. Cheaper, and works
as well for transformers.
"""

import torch
import torch.nn as nn


class AbyssRMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    
    def forward(self, x):
        # normalize in float 32 for stability, then cast back to x's dtype
        dtype = x.dtype
        x = x.float()
        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (self.weight * (x * rms)).to(dtype)