ML Katas

Implement Layer Normalization from scratch in JAX. Layer Normalization is a key component in many modern neural network architectures, like Transformers. It normalizes the inputs across the features.

Formula: $y=\frac{x-E[x]}{\sqrt{Var[x]+ϵ}}*\gamma +\beta$

Where $\gamma$ and $\beta$ are learnable parameters (gain and bias).

Verification: - The output of your layer norm function should have a mean close to 0 and a standard deviation close to 1 before applying the gain and bias. - The shape of the output should be the same as the input.