LSTM Cell From Scratch in NumPy (PyTorch Gate Order)

Compute one step of an LSTM. Inputs: x (length d), h and c (length H), W (4H × (d + H)) and b (length 4H). With z = W·[x, h] + b, the four blocks of H rows are, in order, the input gate i, the forget gate f, the candidate g and the output gate o:

i = σ(z₀), f = σ(z₁), g = tanh(z₂), o = σ(z₃);
c′ = f ⊙ c + i ⊙ g and h′ = o ⊙ tanh(c′).

Return [h′, c′]. Pre-activations can be large in either direction; the sigmoid must not overflow.

Examples

Input x = [1], h = [0], c = [0], W = [[0.5, 0], [0.5, 0], [1, 0], [1, 0]], b = [0, 0, 0, 0]

Output [[0.3227], [0.4741]]

d = H = 1, starting from zero state.

Input x = [0], h = [1], c = [2], W = [[0, 0], [0, 50], [0, 0], [0, 50]], b = [-50, 0, 0, 0]

Output [[0.964], [2]]

Forget gate saturated open, input gate shut: the cell state carries over unchanged.

Constraints

1 ≤ d, H ≤ 64
Pre-activations may reach ±800
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

One matrix product on the concatenation [x, h] gives all four blocks; slice it by H.

Hint 2

Compute c′ first; h′ uses tanh(c′), not tanh of the old c.

Hint 3

For σ, branch on the sign: 1/(1 + e^{−x}) for x ≥ 0 and e^{x}/(1 + e^{x}) for x < 0.

Where candidates go wrong

Using the gate order of another library or paper (forget gate first, for instance).
np.exp(-z) for large negative z: an overflow warning and, in some code paths, nan.
h′ = o ⊙ tanh(c) with the old cell state.

The idea

Concatenate x and h, apply one affine map, split the result into four equal blocks, activate each block with its own function, then update c and h. Each step is elementwise except the one matrix product, so the cell costs O(H(d + H)). The stable sigmoid uses np.where over the two branches, exponentiating only non-positive values.

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])                     # z < 0, so no overflow
    out[~pos] = e / (1.0 + e)
    return out

def lstm_cell(x, h, c, W, b):
    x, h, c = (np.asarray(v, dtype=float) for v in (x, h, c))
    W, b = np.asarray(W, dtype=float), np.asarray(b, dtype=float)
    H = h.shape[0]
    z = W @ np.concatenate([x, h]) + b
    i = sigmoid(z[:H])
    f = sigmoid(z[H:2 * H])
    g = np.tanh(z[2 * H:3 * H])
    o = sigmoid(z[3 * H:])
    c_new = f * c + i * g
    h_new = o * np.tanh(c_new)
    return [h_new.tolist(), c_new.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

The standard check that a candidate can turn a diagram into exact code: slicing one matrix product into four gates in the right order, and computing the new cell state before the new hidden state.

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

Keep going

Medium15 min

One step of an LSTM cell