GRU Cell From Scratch in NumPy, Matching PyTorch

Compute one GRU step in PyTorch's convention. Inputs: x (length d), h (length H), W_ih (3H × d), W_hh (3H × H), b_ih and b_hh (length 3H). The three blocks of H rows are, in order, reset r, update z and new n. With a = W_ih·x + b_ih and u = W_hh·h + b_hh:

r = σ(a₀ + u₀), z = σ(a₁ + u₁), n = tanh(a₂ + r ⊙ u₂);
h′ = (1 − z) ⊙ n + z ⊙ h.

Return h′ as a list.

Examples

Input x = [1], h = [0], W_ih = [[1], [1], [1]], W_hh = [[1], [1], [1]], b_ih = [0, 0, 0], b_hh = [0, 0, 0]

Output [0.2048]

Zero state: h′ is (1 − z) times the candidate.

Input x = [0], h = [1], W_ih = [[0], [0], [0]], W_hh = [[0], [50], [0]], b_ih = [0, 0, 0], b_hh = [0, 0, 0]

Output [1]

Update gate saturated at 1: the state is kept.

Constraints

1 ≤ d, H ≤ 64
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Compute the two affine maps separately; the reset gate needs the hidden one on its own.

Hint 2

r multiplies (W_hn h + b_hn), the third block of the hidden map, after the matrix product.

Hint 3

z weights the old state, not the candidate: h′ = (1 − z) n + z h.

Where candidates go wrong

Applying r to h before the multiplication, as in the original paper: a valid GRU, but not PyTorch's.
Swapping the roles of z and 1 − z.
Adding the two biases before splitting and then forgetting that b_hn sits inside the reset product.

The idea

Keep the input map a = W_ih x + b_ih and the hidden map u = W_hh h + b_hh separate. The reset and update gates use the sums of their blocks; the candidate adds the input block to the reset-scaled hidden block. The final state interpolates between the candidate and the old state.

Time O(H · (d + H)) · Space O(H)

Reference solution

import numpy as np

def sigmoid(z):
    out = np.empty_like(z)
    pos = z >= 0
    out[pos] = 1.0 / (1.0 + np.exp(-z[pos]))
    e = np.exp(z[~pos])
    out[~pos] = e / (1.0 + e)
    return out

def gru_cell(x, h, W_ih, W_hh, b_ih, b_hh):
    x, h = np.asarray(x, dtype=float), np.asarray(h, dtype=float)
    H = h.shape[0]
    a = np.asarray(W_ih, dtype=float) @ x + np.asarray(b_ih, dtype=float)
    u = np.asarray(W_hh, dtype=float) @ h + np.asarray(b_hh, dtype=float)
    r = sigmoid(a[:H] + u[:H])
    z = sigmoid(a[H:2 * H] + u[H:2 * H])
    n = np.tanh(a[2 * H:] + r * u[2 * H:])          # reset applied after the product
    return ((1 - z) * n + z * h).tolist()

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

"Implement a GRU that matches the library" separates candidates who have read the equations from those who have read the documentation. The reset gate's position is the whole question.

The theory behind it is in the lesson LSTM and GRU cells, and beam search, from scratch.

Keep going

Medium15 min

One step of a GRU cell, as PyTorch computes it