import torch
import torch.nn as nn


class LayerNorm(nn.Module):
    """
    Layer normalisation: squeezes outputs to mean 0 and variance 1.
    eps avoids division by zero.
    scale/shift are learnable affine parameters.
    """

    def __init__(self, embed_dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.embed_dim = embed_dim
        self.scale = nn.Parameter(torch.ones(self.embed_dim))
        self.shift = nn.Parameter(torch.zeros(self.embed_dim))

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