src/kernels/normalization/layer_norm.py

"""
LayerNorm (Layer Normalization).

The original normalization technique for transformers (Ba et al., 2016).
LayerNorm normalises each token's hidden state by subtracting the mean and
dividing by the standard deviation, then applies a learnable affine
transform:  output = scale * (x - mean) / sqrt(var + eps) + shift.

Compared to RMSNorm (which skips mean-centering and the shift parameter),
LayerNorm has slightly higher compute cost due to the extra mean subtraction
and additive bias.  In practice, modern LLMs (LLaMA, Qwen3, Gemma) prefer
RMSNorm because the mean-centering provides negligible benefit in pre-norm
architectures, while the savings add up across dozens of layers.

Included here for educational comparison with RMSNorm.

Computation is done in float32 regardless of input dtype to avoid
numerical instability in half-precision, then cast back.
"""

import torch
import torch.nn as nn


class LayerNorm(nn.Module):
    def __init__(self, emb_dim: int, eps: float = 1e-5, dtype=None):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(emb_dim, dtype=dtype))
        self.shift = nn.Parameter(torch.zeros(emb_dim, dtype=dtype))

    def forward(self, x):
        input_dtype = x.dtype
        # Force float32 for numerical stability (matches RMSNorm convention)
        x = x.to(torch.float32)
        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.float() * norm_x + self.shift.float()).to(input_dtype)