Foundations · INTEGRATION

LayerNorm and RMSNorm from reductions

MediumRaw PyTorchAbout 20 min3 public examples plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Rescale each final-axis feature vector using one shared denominator: LayerNorm first subtracts its mean and uses its centered mean square; RMSNorm uses the original mean square. Apply the supplied scale and optional bias to each channel, returning the input shape.

Function signature

layer_rms_normalization(
    x: torch.Tensor,
    scale: torch.Tensor,
    bias: torch.Tensor | None = None,
    variant: str = 'layer',
    eps: float = 1e-05,
) -> torch.Tensor

Worked example

A varying row and a constant row

For [1,3], the mean is 2 and population variance is 1. With eps=1, normalize to $[-1,1]/\sqrt2$ , then multiply by [2,-1] and add [0.5,1], giving approximately [-0.914214,0.292893]. The constant row [4,4] centers to zero and returns the bias [0.5,1].

import torch

x = torch.tensor([[1.0, 3.0], [4.0, 4.0]], dtype=torch.float64, requires_grad=True)
scale = torch.tensor([2.0, -1.0], dtype=torch.float64, requires_grad=True)
bias = torch.tensor([0.5, 1.0], dtype=torch.float64, requires_grad=True)
variant = 'layer'
eps = 1.0

result = layer_rms_normalization(x=x, scale=scale, bias=bias, variant=variant, eps=eps)

Expected return

torch.tensor([[-0.9142135623730949, 0.29289321881345254], [0.5, 1.0]], dtype=torch.float64)

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[\ldots, C]$ Feature vectors to normalize. Leading axes may be empty; the final width must be positive.
`scale`torch.Tensor	$[C]$ One supplied multiplier per feature.
`bias`torch.Tensor \| NoneDefault `None`	$[C]$ Optional offset per feature; omission means zero.
`variant`strDefault `'layer'`	Choose "layer" for centered population-variance normalization or "rms" for uncentered mean-square normalization. Both variants support optional bias in this teaching contract.
`eps`floatDefault `1e-05`	Strictly positive value added to the population variance (LayerNorm) or mean square (RMSNorm), inside the square root.

Return

torch.Tensor · Exactly x.shape.

Return Floating normalized-and-affine tensor with exactly the shape, dtype and device of x.

Required behavior

Normalize only the last axis of $x$ , with width $C\ge1$ . Every combination of leading indices defines an independent feature vector.
LayerNorm: $\mu=\frac1C\sum_c x_c$ , $v=\frac1C\sum_c(x_c-\mu)^2$ , and $z_c=(x_c-\mu)/\sqrt{v+\varepsilon}$ . Use population variance, not division by $C-1$ .
RMSNorm: $z_c=x_c/\sqrt{\frac1C\sum_c x_c^2+\varepsilon}$ . Do not subtract a mean.
In either variant, return $y=z\odot\mathrm{scale}+\mathrm{bias}$ . The scale vector is supplied and required. Missing bias means zero; optional RMSNorm bias is an explicit teaching extension.
eps is finite and strictly positive, added inside the square root after the channel mean. All floating tensors are finite and share dtype/device.
A constant LayerNorm row, including a single-channel row, returns bias exactly. An all-zero RMSNorm row does the same. Input gradients remain finite.
Empty leading axes are allowed and retained; an empty channel axis is not. Preserve noncontiguous inputs without mutation. There is no mask, running state, dropout, or residual addition.

Why this matters

LayerNorm removes a row’s mean before scaling its spread; RMSNorm scales its magnitude without removing the mean.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def layer_rms_normalization(
    x: torch.Tensor,  # (..., C)
    scale: torch.Tensor,  # (C,)
    bias: torch.Tensor | None = None,  # (C,)
    variant: str = 'layer',
    eps: float = 1e-05,
) -> torch.Tensor:
    """Return a tensor with shape x.shape."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement layer_rms_normalization")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.