10. Minimal GRU Cell From Scratch

Problem

Every gated-recurrence family used in efficient sequence models — from the classic GRU/LSTM line through today's data-controlled linear recurrences (GateLoop, Mamba's selective SSM) — shares one idea: a per-step gate decides how much of the previous hidden state to keep versus overwrite with new information. Before reasoning about any of the more exotic gating schemes an interviewer might bring up, you should be able to write the canonical GRU update rule directly against raw tensors, with the gate math explicit rather than hidden behind nn.GRUCell.

The reset gate r_t controls how much of the past state leaks into the new candidate; the update gate z_t then interpolates between fully committing to that fresh candidate and fully preserving the old state — exactly the "keep vs. overwrite" mechanism that generalizes into the data-controlled decay terms used by modern SSM-style layers.

Source: src/10_gru_cell_from_scratch.py

def gru_cell(
    x: torch.Tensor, h_prev: torch.Tensor,
    weight_ih: torch.Tensor, weight_hh: torch.Tensor,
    bias_ih: torch.Tensor, bias_hh: torch.Tensor,
) -> torch.Tensor: ...

>>> cell = torch.nn.GRUCell(4, 3)
>>> x, h0 = torch.randn(2, 4), torch.randn(2, 3)
>>> h1 = gru_cell(x, h0, cell.weight_ih, cell.weight_hh, cell.bias_ih, cell.bias_hh)
>>> h1.shape
torch.Size([2, 3])

Step-by-Step Approach

1. Compute the two stacked pre-activations: gi = x @ weight_ih.T + bias_ih and gh = h_prev @ weight_hh.T + bias_hh, each shape (batch, 3*hidden).
2. Split each into three equal chunks along the last dim, in the order [reset, update, new] — this must match nn.GRUCell's own row ordering in weight_ih/weight_hh exactly, or the gates get crossed and the output is silently wrong.
3. Compute the reset gate r = sigmoid(i_r + h_r) and update gate z = sigmoid(i_z + h_z).
4. Compute the candidate state n = tanh(i_n + r * h_n) — note r multiplies only the hidden-state contribution to the candidate, never the input contribution.
5. Interpolate: h_new = (1 - z) * n + z * h_prev — z close to 1 keeps the old state, z close to 0 fully commits to the new candidate.
6. Sanity-check against torch.nn.GRUCell with the same weights copied in — a correct from-scratch implementation should reproduce it to floating-point precision.

The key insight is that a GRU update is just two matmuls, three gate splits, and one convex interpolation — the reset gate is only ever applied to the hidden-state branch of the candidate, which is the detail most implementations get wrong from memory.

Reference solution

import torch


def gru_cell(
    x: torch.Tensor,
    h_prev: torch.Tensor,
    weight_ih: torch.Tensor,
    weight_hh: torch.Tensor,
    bias_ih: torch.Tensor,
    bias_hh: torch.Tensor,
) -> torch.Tensor:
    # (batch, input_size) @ (input_size, 3*hidden) -> (batch, 3*hidden)
    gi = x @ weight_ih.T + bias_ih
    gh = h_prev @ weight_hh.T + bias_hh
    # split the stacked [reset, update, new] blocks -- must match nn.GRUCell's
    # row ordering in weight_ih/weight_hh exactly, or the gates get crossed
    i_r, i_z, i_n = gi.chunk(3, dim=-1)
    h_r, h_z, h_n = gh.chunk(3, dim=-1)
    r = torch.sigmoid(i_r + h_r)
    z = torch.sigmoid(i_z + h_z)
    # reset gate is applied only to the hidden-state contribution to the
    # candidate, never to the input contribution
    n = torch.tanh(i_n + r * h_n)
    # z close to 1 -> keep old state; z close to 0 -> fully commit to candidate
    h_new = (1.0 - z) * n + z * h_prev
    return h_new

Key Functions & Tricks

• tensor.chunk(3, dim=-1) — splits the stacked 3*hidden pre-activation into the three gate blocks in one call, matching nn.GRUCell's packed weight layout.
• torch.sigmoid — squashes both gates into (0, 1) so they behave as "how much" fractions.
• torch.tanh — keeps the candidate state bounded in (-1, 1), preventing unbounded growth across many recurrent steps.
• weight_ih.T / weight_hh.T — nn.GRUCell stores weights as (out_features, in_features), so the matmul needs the transpose to line up dimensions.
• torch.nn.GRUCell — used here purely as ground truth (weights copied out and fed into the from-scratch function), not called in the solution path itself.

How to Recognize This Pattern