import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    """RMSNorm (Section 4.7): cheaper alternative to LayerNorm.

    Rescales by root-mean-square of the activations instead of full
    mean/variance normalization. No bias, single learnable scale per dim.
    """

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # compute in float32 for stability regardless of input dtype (bf16 etc.)
        dtype = x.dtype
        x = x.float()
        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        out = x * rms
        return (out.to(dtype)) * self.weight