This is the shape a pytorch training script has, written as code with nothing hidden in it: a batch of sequences trains a two-block pre-norm transformer, and every mechanism — the embedding lookup, the batched attention, the dropout, the schedule, the optimizer — is a line you can point at.

The example is examples/BatchedTrainer in the tutorial repository.

Getting the file

make -C examples/BatchedTrainer run
step 1: a = 0.062500 n = 1.762002 loss = 2.451075
step 2: a = 0.046875 n = 0.689282 loss = 1.674839
step 3: a = 0.015625 n = 0.681793 loss = 1.709891
e0 = -0.373581 -0.278466 wq0 = -0.655883 g1 = 1.078082

Three training steps over a batch of two three-token sequences: the schedule’s alpha, the embedding gradient’s norm before clipping, and the loss — then a few trained weights. Every digit has been checked against an independent oracle that replays the entire computation, and every digit repeats at every optimization tier and under any carrier count. The gelu, the cross-entropy and the cosine schedule call the C library’s tanh, exp, log and cos, so the last digits are the ones the machine’s C library gives.

The two shapes of a batch

Almost everything a transformer does is per token: the projections, the normalizations, the activation, the loss. Those compute over one flat matrix of six token rows — matrix[6, 4]. Only the attention score product needs to know which rows form a sequence — pytorch writes x.view(B, T, C) there, and Mica writes Reshape:

q3 := Reshape(q);              { matrix[6, 4] read as tensor[2, 3, 4] - same cells }
sc := q3 * Transpose(k3);      { the batched score product, [2, 3, 3] }
...
o := Reshape(o3);              { and back to the flat per-token matrix }

Reshape re-reads the same row-major cells in the assignment target’s static shape; the cell counts are proven equal at compile time, and on tracked values the reshape records, its derivative copying the gradient back cell for cell. The shape comes from where the call stands — assign it directly to a shaped variable, and the left side names the new shape.

One block, every line visible

a := RmsNorm(x, blocks[i].g1);
q := a * blocks[i].wq;  k := a * blocks[i].wk;  v := a * blocks[i].wv;
q3 := Reshape(q);  k3 := Reshape(k);  v3 := Reshape(v);
sc := q3 * Transpose(k3);
sc := 0.5 * sc;                { 1/sqrt(C), exactly representable }
sc := sc + m3;                 { m3 := SliceSpread(cm) - the mask per batch slice }
at := RowSoftmax(sc);
at := Dropout(at, dm3);        { the seeded mask, drawn this step }
o3 := at * v3;
o := Reshape(o3);
x := x + o * blocks[i].wo;
a := RmsNorm(x, blocks[i].g2);
h := GeluMap(a * blocks[i].w1);
x := x + h * blocks[i].w2;

The row-wise verbs — RmsNorm, RowSoftmax, the activation maps, CrossEntropy — take the batched shapes without new names: a higher-rank operand enters as its trailing-axis matrix, every leading axis folded into the row count. RowSoftmax over [2, 3, 3] is six rows of three, exactly the reading its mathematics has.

The loss and the walk

lg := f * Transpose(e);            { the tied head: logits [6, 5] }
loss := CrossEntropy(lg, ti);      { targets as a matrix[2, 3] of int64 }

The targets arrive as a matrix — one class index per token, in the batch’s own shape. After Backward, the walk is the loop you read about in Model depth, at full size:

gm := Gradient(e);
nrm := ClipGradNorm(address gm, 1.0);
AdamWStep(address ep, gm, address me, address ve, step, alpha, 0.9, 0.999, 0.00000001, 0.0625);

The embedding’s gradient collects two roads — the Gather lookup scattering rows back, and the tied head flowing through the transpose — is clipped, and steps under the schedule’s alpha. Then every block’s eight parameters take the same three sentences each. At step one the printed norm is 1.762002 — above the ceiling, so the clip fired; from step two it stays under.

Gradient accumulation, when you want it

There is no machinery for it because none is needed: gradients extracted with Gradient are plain values, and accumulating across windows is arithmetic —

gsum := gsum + Gradient(w);

— record a window per micro-batch, add the extracted gradients, and step once with the sum. What pytorch does behind loss.backward()’s accumulation semantics is a + you write.

What stays compile-checked

The batch size, context length, and channel widths are constants, so every shape in the program is proven at compile time — the score product cannot meet the wrong axis, the reshape cannot miscount, and a wrong gain length is a compile error, not a runtime surprise. A vocabulary read from data at run time is its own story: the dynamic tensors family carries runtime extents through the operators and Gather, and the recorded row-wise verbs hold their static-extents rule — the honest boundary between shapes in types and shapes in data, stated by the compiler when you cross it.