LSTMwith its shapes on

I · The problem

What is h?

An isolated equation tells you less than an implementation needs. A paper may define the missing information in its prose; we want it visible at every important use.

Pedagogical modernisation of the LSTM family, not verbatim paper equations: batching is implicit here, peepholes are omitted, the candidate is named g, and ⊙ distinguishes element-wise multiplication. Compare Graves 2013 §2.1, Eqs. 7–11 and the source-linked translations.

Questions to ask before writing the code

  1. What is h? One sequence has 768 hidden coordinates; a training batch may have 16×768, and all times together give 16×128×768. Vector equations commonly describe one example, not the entire batched implementation.
  2. Which way round are the weights? Column notation uses W·x; row-oriented batched code often uses x@W. PyTorch stores a transposed layout relative to the row-oriented W here. A symbol alone does not state the convention.
  3. Which product is meant? Some papers distinguish matrix and element-wise multiplication typographically; others rely on prose. We insist on a dot for contraction and ⊙ for element-wise work.
  4. Which axes broadcast? A bias of 768 entries is repeated down 16 rows, not learned separately for each sequence. Explicit broadcasting makes the parameter count checkable.
  5. Where are the four gates? An implementation may store all four projections in one fused output of width 3072. Its gate order must be declared; this site uses i,f,g,o.
  6. Which “768” is the source? A 768×768 dense map has two different coordinate roles. Equal sizes do not justify naming both axes D.

The shapes are not extra commentary. They are part of the computation.

Old and new, one gate

Use B=16, E=512, D=768. The output coordinates R also have size 768, but are a different axis identity from recurrent input coordinates D.

Shape-implicit modernised form
Ledger presentation

The first line projects into R. The second explicitly renames R→D by matching equal-numbered coordinates, so the gate can multiply the cell state element-wise. There is no D×D weight that would leave a named contraction ambiguous.

Fused, as an implementation usually stores it

K=G·R has 4×768=3072 entries and is distinct from D. Unpack unfuses K into G×R, selects the gates, then renames R→D. The full specification states these rules and acknowledges their named-tensor antecedents.

TensorAxesStored entries
input at one stepB×E = 16×5128,192
state h or cB×D = 16×76812,288 each
input projection WE×K = 512×30721,572,864
recurrent projection UD×K = 768×30722,359,296
stored bias bK = 3072, virtual B at use3,072

4D(E+D+1)=3,935,232 trainable parameters, independent of B and T. Naming dimensions lets us check that claim directly.