(B, C, H, W) and (B, H, T, D) are the same idea: a flat list of numbers plus a shape tuple and strides. Every reshape, transpose and broadcast is bookkeeping — and every shape error is a mismatch you can settle with pen and paper.
A tensor is a flat buffer of numbers plus a shape tuple. Rank counts the axes, and the element count is their product: (2, 3, 4) holds 24 numbers. Wrap your head around the tuple and every framework error becomes readable.
Stride says how many positions to skip per axis. A contiguous tensor lays its logical rows down in order; a transpose swaps strides without moving data, so neighbours stop being neighbours and view() refuses. reshape() copies when it must.
strideᵢ = product of later sizes03 / ONE NOTATION FOR EVERY CONTRACTION
Einsum: keep the output letters, sum the rest.
Label each axis with a letter. Letters that survive into the output are free; letters left out are multiplied and summed. Dot product, matmul, outer product, trace, transpose and attention all become one readable line — with the shape written into it.
“bhtd,bhsd->bhts”
MENTAL MODEL IN ONE SENTENCE
A tensor is a flat buffer plus shape and strides; every operation you meet — reshape, transpose, broadcast, reduce, einsum — is bookkeeping on that metadata, and every shape error is a mismatch you can settle with pen and paper.
By the end you will be able to read any tensor shape on sight, predict broadcast results at any rank, explain when view() works and when reshape() must copy, write and read einsum expressions, and trace the exact shapes through every step of multi-head attention.
01
ONE BUFFER, MANY AXES
A tensor is a grid with any number of axes.
A scalar, a vector, a matrix, an image batch: the same idea at four ranks. The shape tuple says how much lives along each axis, and the strides say where to find it.
A tensor is a multi-dimensional array of numbers with a uniform type. The number of axes is its rank (also called order), each axis is a dimension, and the shape is the tuple that lists the size along each axis. A scalar has shape () and rank 0, a vector of three numbers is (3,) with rank 1, a 2 × 3 matrix is (2, 3) with rank 2, and a 3D tensor might be (2, 3, 4) with rank 3 — 2 · 3 · 4 = 24 numbers.
Plain English: the shape answers “how many along each direction”, and the element count is the product of those answers. Indexing works the way you would guess: x[1, 2, 3] picks one number, an integer index drops that axis, and a slice keeps it.
rank 0 scalar () 1 number
rank 1 vector (3,) 3 numbers
rank 2 matrix (2, 3) 6 numbers
rank 3 tensor (2, 3, 4) 2·3·4 = 24 numbers
shapes you will meet daily element count
images (PyTorch) (B, C, H, W) = (32, 3, 224, 224) 4,816,896
images (TensorFlow) (B, H, W, C) = (32, 224, 224, 3) 4,816,896 ← same
tokens (B, T, D) = (16, 128, 768) 1,572,864
attention (B, H, T, D) = (16, 12, 128, 64) 1,572,864
linear weight (out, in) = (10, 784) 7,840
indexing a (3, 4) tensor x:
x[:, 0] fix axis 1, drop it → (3,)
x[0] fix axis 0, drop it → (4,)
x[0:2] keep axis 0 → (2, 4)
Shape workbench
Add and remove axes, change each size, and reorder the axes. The shape tuple and strides update; the block is the same numbers rearranged.
shape (2, 3, 4) = 2 · 3 · 4 = 24 numbers
strides (12, 4, 1)
last element [1][2][3] → flat 23
= 1·12 + 2·4 + 3·1 = 23
(0-based, so this is element number 23 of 24)
Reordering axes never moves a number: the strides tuple reorders with them. That is why transpose is free and why it breaks contiguity.
Derivation: strides — how a grid lives in one flat line
Memory is one-dimensional; a tensor is a flat buffer plus a recipe for reading it. Row-major (“C order”) stores the last axis fastest, so the stride of an axis is the product of every size to its right.
Shape (3, 4): moving one step along axis 1 jumps 1 position; moving one step along axis 0 must skip a whole row of 4. Strides (4, 1).
In general stride[i] = shape[i+1] · shape[i+2] · …. Shape (2, 3, 4) → strides (12, 4, 1).
The flat position of an index is the dot product of the index with the strides: Σ index[i] · stride[i].
shape (3, 4), strides (4, 1)
flat = [a b c d | e f g h | i j k l]
element [1][2] → 1·4 + 2·1 = 6 → g ✓
shape (2, 3, 4), strides (12, 4, 1)
element [1][2][3] → 1·12 + 2·4 + 3·1 = 23
the 24th and last number of the buffer ✓
Every operation in this lesson is a story about this formula. Reshape rewrites the strides, transpose swaps them, broadcasting reads a stride-0 axis — and none of them move the underlying numbers.
Quick check
A tensor has shape (2, 3, 4). How many axes does it have, and how many numbers does it hold?
02
BROADCASTING, THE FULL RULE
Align from the right. Stretch the ones.
Lesson 02 met broadcasting for a bias vector. The same rule works at any rank: pad the shorter shape with 1s on the left, then compare each axis — equal or 1 passes, anything else errors.
Broadcasting lets two tensors of different shapes combine without copying either one. Write both shapes right-aligned; a missing leading axis counts as size 1; an axis of size 1 is re-read for every index along that axis. The result takes the larger size at each position. That is the whole algorithm — the rest is bookkeeping.
unsqueeze exists to make the alignment line up on purpose. A bias (D,) added to a batch (B, T, D) already aligns with the trailing axis; but a per-channel image scale (C,) applied to (B, C, H, W) must be reshaped to (1, C, 1, 1) first, or it will be compared against W.
A (8, 1, 6, 1)
B (7, 1, 5) ← pad with 1s on the left
B (1, 7, 1, 5)
axis −4: 8 vs 1 → 8 B re-read along this axis
axis −3: 1 vs 7 → 7 A re-read
axis −2: 6 vs 1 → 6 B re-read
axis −1: 1 vs 5 → 5 A re-read
result (8, 7, 6, 5) = 1680
A stores 8·6 = 48 numbers, B stores 7·5 = 35.
No axis has a non-1 on both sides, so every A value meets
every B value exactly once:
48 · 35 = 1680 = 8·7·6·5 ✓
Broadcasting in three dimensions and up
Shapes align from the right. A 1 is a ghost: it is re-read for every value along that axis, never copied. Equal sizes pass through.
A
8161
8 · → 7 · 6 · → 5
B
1715
→ 8 · 7 · → 6 · 5
OUT8765
A (8, 1, 6, 1)
B (7, 1, 5)
pad B (1, 7, 1, 5)
axis 0: 8 vs 1 → 8 (B stretched)
axis 1: 1 vs 7 → 7 (A stretched)
axis 2: 6 vs 1 → 6 (B stretched)
axis 3: 1 vs 5 → 5 (A stretched)
→ (8, 7, 6, 5)
A has 48 numbers, B has 35, result has 1680
ghost reads (never stored):
A re-read: axis 1 ×7, axis 3 ×5
B re-read: axis 0 ×8, axis 2 ×6
the flagship example: every non-1 axis comes from a different side
A size-1 axis has one stored value, so it can answer for any index — that is the whole trick. Missing leading axes are treated as size 1.
The rule, step by step, with numeric checks
Write both shapes right-aligned and pad the shorter one with 1s:(2, 3) and (3,) become (2, 3) and (1, 3).
Compare axis by axis. Compatible when the sizes are equal, or one of them is 1. The result takes the max.
A size-1 axis is virtually repeated: its single value answers for every index. No memory is shared or copied — the read is simply allowed.
If a pair is incompatible, raise a shape error and stop.
(2, 3) + (3,) → (2, 3) each row gets the same 3-bias
(2, 3) + (2, 1) → (2, 3) one bias per row, stretched across columns
(2, 3) + (2,) → 3 vs 2: ERROR (this is the Lesson 02 column trap)
(B, T, D) + (D,) → (B, T, D) bias per feature
(B, C, H, W) * (C,) → ERROR align: C vs W
(B, C, H, W) * (1, C, 1, 1) → (B, C, H, W) the fix
numeric outer check, (3, 1) + (1, 4):
[[1], [2], [3]] + [[10, 20, 30, 40]]
= [[11, 21, 31, 41],
[12, 22, 32, 42],
[13, 23, 33, 43]] every one of the 12 sums reads one A and one B value
Quick check
What is the result shape of (4, 1, 3) + (2, 3)?
03
ELEMENT-WISE OPS & REDUCTIONS
Same shape, or one axis fewer.
Element-wise operations touch every position and keep the shape. Reductions collapse the axes you name — and the axis you name is the whole game.
Element-wise operations (add, multiply, subtract, ReLU, exp) apply independently at every position: one number in, one number out, shape unchanged. Reductions (sum, mean, max) collapse one or more axes to a single value, removing them from the shape. A reduction over axis=0 removes the outermost axis; axis=1 removes the next one. That choice — which axis to collapse — is the single most common source of silent math bugs.
Plain English:sum(axis=0) means “add up the rows, one column at a time”; sum(axis=1) means “add up the columns, one row at a time”. keepdims=True leaves the collapsed axis behind as a size-1 axis, which is exactly what you need to broadcast the result back.
A reduction names the axes to collapse. Global average pooling turns (B, C, H, W) into (B, C); token mean pooling turns (B, T, D) into (B, D). The numbers being averaged are re-read, then replaced by their mean.
A = [[1, 2],
[3, 4]]
element-wise A * A = [[1, 4], [9, 16]] shape stays (2, 2)
sum(axis=0) [1+3, 2+4] = [4, 6] collapse rows → (2,)
sum(axis=1) [1+2, 3+4] = [3, 7] collapse columns → (2,)
sum() 10 collapse every axis → scalar
mean(axis=0) [2, 3]
softmax of logits [2, 1, 0] over the last axis:
exp = [7.389, 2.718, 1.000]
sum = 11.107
out = [0.665, 0.245, 0.090] sums to 1.000 ✓ shape unchanged
(B, C, H, W).mean(axis=[2, 3]) → (B, C) global average pooling
(B, T, D).mean(axis=1) → (B, D) sequence mean pooling
(B, T, D).mean(axis=1, keepdims=True) → (B, 1, D) broadcasts back over T
04
RESHAPE, VIEW, PERMUTE
Same numbers. New shape — maybe new memory.
Reshape rewrites the shape metadata. Transpose and permute reorder the axes. One of those operations is free, one can silently copy, and one can throw.
Reshape keeps the flat element order and gives it a new shape; the element count must match, and -1 lets one dimension be inferred. Transpose swaps two axes and permute reorders all of them; neither moves a number, they just swap the strides. View is a reshape that promises to reuse the same buffer — and that promise can fail.
Plain English: a tensor is a flat list plus a reading recipe. Reshape changes the recipe’s grouping, transpose changes the recipe’s order, and view insists the old buffer can satisfy the new recipe. After a transpose it usually cannot, because the numbers a logical row needs are now spread apart in memory.
reshape (2, 3, 4) → (6, 4): same flat order, same 24 numbers
(2, 3, 4).reshape(-1, 3) → (8, 3) because 24 / 3 = 8
transpose (3, 4) → (4, 3):
before shape (3, 4) strides (4, 1) contiguous
after shape (4, 3) strides (1, 4) non-contiguous
logical row 0 of the transpose reads buffer slots 0, 4, 8 — not 0, 1, 2.
permute NCHW → NHWC, on (1, 2, 3, 4):
permute(0, 2, 3, 1) → shape (1, 3, 4, 2), strides (24, 4, 1, 12)
same 24 numbers; only the stride tuple changed.
Reshape vs view: follow the buffer
Transpose swaps the strides and moves nothing. Watch the connecting lines cross, then materialize a copy and watch them straighten.
shape (3, 4)
strides (4, 1)
contiguous? yes — stride[i] = product of later sizes
buffer the same 12 numbers, untouched
cell [0][1] holds 1
run .view or .reshape to compare
view() is a promise: “same buffer, new shape”. reshape() keeps the promise when it can and quietly copies when it cannot.
Permute: axes change places, numbers stay put
Move the axis bars, or jump to a conversion preset. The shape and strides tuples reorder together; the buffer never moves.
permute is the general form of transpose: one call can reorder every axis at once, which is exactly how NCHW becomes NHWC.
Derivation: the contiguity test, and why view refuses
A tensor is contiguous when logical neighbours are memory neighbours. The test is one line: read the axes left to right and check stride[i] = product of all sizes to its right.
Why the rule works: rows of the logical grid are laid down one after another. If an axis’s stride is smaller than the block it must skip, its values are interleaved with another axis’s.
view() refuses non-contiguous input because it will not guess how to regroup the buffer. reshape() checks the same condition and, when it fails, copies the data into a contiguous buffer first.
transpose of (3, 4) with buffer [0 1 2 3 | 4 5 6 7 | 8 9 10 11]:
transpose row 0 = [0, 4, 8] buffer slots 0, 4, 8
view(6, 2) would assume row 0 = slots 0, 1 and row 1 = slots 2, 3.
Those slots hold [0, 1] and [2, 3] — not the transposed rows.
PyTorch raises instead of quietly returning wrong numbers.
after .contiguous(): buffer [0 4 8 1 5 9 2 6 10 3 7 11]
shape (4, 3), strides (3, 1): 3 = 3 ✓, 1 = 1 ✓ → view(6, 2) works
(and shape (4, 3) is preserved element by element: 4 is still row 0, col 1.)
Quick check
A (3, 4) tensor is transposed to (4, 3). Its strides change from (4, 1) to (1, 4). What does that mean?
05
NCHW VS NHWC
Two layouts, one silent disaster.
The same image batch can be stored channels-first or channels-last. PyTorch and TensorFlow disagree by default, and the wrong pairing produces garbage rather than an error.
A batch of images has four axes: batch, channels, height, width. The only question is their order. PyTorch defaults to (B, C, H, W) — channels first — where each channel is a complete image plane. TensorFlow, Keras and most mobile runtimes default to (B, H, W, C) — channels last — where every pixel carries its channel values together. NCHW and NHWC are the same tensor with the axes permuted.
Plain English: the layout decides which axis changes fastest in memory. Convolution kernels are written for one layout, so converting a model between frameworks means permuting the axes and making the result contiguous — and doing it once at the boundary, not inside every layer.
Schematic, not to scale. Channels-first keeps each channel as a full image plane; channels-last interleaves the channel values inside every pixel position. Convolution kernels prefer one order or the other — an operation that expects the wrong one does not fail loudly, it reads the right numbers as the wrong axes.
Framework / op
Layout
Typical shape
PyTorch image batch
NCHW
(32, 3, 224, 224)
TensorFlow / Keras image batch
NHWC
(32, 224, 224, 3)
PyTorch Conv2d weight
(out_c, in_c, kH, kW)
(64, 3, 3, 3)
TensorFlow Conv2D weight
(kH, kW, in_c, out_c)
(3, 3, 3, 64)
convert, smallest workable example (2, 3, 4, 5):
x.permute(0, 2, 3, 1) → shape (2, 4, 5, 3), strides (60, 5, 1, 20)
before: strides (60, 20, 5, 1) — the same 120 numbers, re-ordered read
PyTorch can keep NCHW shapes with channels-last memory:
x = x.to(memory_format=torch.channels_last)
x.shape → still (B, C, H, W)
x.stride() → channels now innermost: the fast axis is C
diagnose before you debug:
x.shape, x.stride(), x.is_contiguous()
06
EINSUM
Label every axis. Sum what you drop.
Einstein summation is one notation for dot products, matmuls, transposes, traces and attention. The rule is a sentence: indices in the output are kept, indices missing from the output are summed.
Write a letter for every axis of every input, then write the letters you want in the output. An index that appears in the inputs but not the output is contracted: the operation multiplies along it and sums. Indices that survive into the output are free, and their output order is exactly the order you write. That is the entire language — and it spells the shape contract of every operation at once.
Plain English: “ik,kj->ij” reads “for every i and j, multiply the matching k entries and add them”. “bhtd,bhsd->bhts” reads “keep batch, head, query and key; for each of them, sum over the head dimension”. Writing the expression is writing the loop, and the shape of the answer is the letters you kept.
Expression
Reads as
Result shape
i,i->
dot product: overlap the axes and sum
scalar
i,j->ij
outer product: keep both axes
(I, J)
ii->
trace: only the diagonal survives
scalar
ij->ji
transpose: relabel and reorder
(J, I)
ik,kj->ij
matrix multiply: sum over k
(I, J)
bij,bjk->bik
batched matmul: one product per b
(B, I, K)
bhtd,bhsd->bhts
attention scores: sum over head dim d
(B, H, T, S)
bhts,bhsd->bhtd
attention output: sum over key axis s
(B, H, T, D)
Einsum explorer
Every letter is an axis. Letters you keep in the output are free; letters you drop get multiplied and summed. Type any expression or start from a preset.
bhtd, bhsd -> bhts
TENSOR 1 · SHAPE (32, 12, 128, 64)
bb=32hh=12tt=128dd=64
TENSOR 2 · SHAPE (32, 12, 128, 64)
bb=32hh=12ss=128dd=64
OUTPUT · SHAPE (32, 12, 128, 128)
bb=32hh=12tt=128ss=128
summed (contracted): d — the loop adds over d=64
one batch, one head, T = S = D = 2
Q = [[1, 0], [0, 1]] K = [[1, 0], [0, 1]]
scores[t][s] = Q[t][0]·K[s][0] + Q[t][1]·K[s][1]
scores = [[1, 0],
[0, 1]] shape (1, 1, 2, 2)
the diagonal is “each token attends to itself”
free axes b, h, t, s → result shape (32, 12, 128, 128)
contracted axes d (multiplied and summed)
input shapes (32, 12, 128, 64) × (32, 12, 128, 64)
cost product of every index size
= b(32)·h(12)·t(128)·d(64)·s(128)
= 402,653,184 multiply-adds
d is contracted: every query token dotted with every key token
dimensions shown are the lesson defaults: b=32, h=12, t=128, s=128, d=64, i=128, j=64, k=128; other letters = 3.
The cost counts every combination of kept and summed indices — which is why attention scores, with two sequence-length indices, grow with T². Press “attention scores” to see the expression from the chapter.
Derivation: einsum is a loop, and the cost is a product
Every expression expands into nested loops: one loop per contracted index, and one output slot per combination of free indices. The cost of a contraction is therefore the product of all index sizes, kept and summed together.
"ik,kj->ij" C_ij = Σ_k A_ik · B_kj k is contracted
"i,i->" s = Σ_i a_i · b_i the dot product
"i,j->ij" M_ij = a_i · b_j nothing summed
"ii->" t = Σ_i A_ii the trace
"ij->ji" B_ji = A_ij pure relabelling
"bij,bjk->bik" C_bik = Σ_j A_bij · B_bjk one matmul per batch
"bhtd,bhsd->bhts" S_bhts = Σ_d Q_bhtd · K_bhsd query · key per head
numeric check — the second worked example:
A (3, 2) = [[1, 2], [3, 4], [5, 6]]
B (2, 3) = [[7, 8, 9], [10, 11, 12]]
C[2][1] = A[2][0]·B[0][1] + A[2][1]·B[1][1]
= 5·8 + 6·11 = 40 + 66 = 106
C = [[27, 30, 33],
[61, 68, 75],
[95, 106, 117]] shape (3, 3) ✓
cost = product of every index size:
"bij,bjk->bik" with b=32, i=128, j=64, k=128
32 · 128 · 64 · 128 = 33,554,432 multiply-adds
"bhtd,bhsd->bhts" with b=2, h=4, t=s=8, d=16
2 · 4 · 8 · 8 · 16 = 8,192 multiply-adds
double the sequence length and the score cost quadruples:
both t and s are free indices of size T.
A kept index can be summed later; a dropped index cannot come back. That is why the one rule needs no special cases: transpose, trace and contraction fall out of the same loop.
Quick check
In “ij,jk->ik”, which letter is contracted, and what is the result shape if i=3, j=2, k=4?
07
ATTENTION, SHAPE BY SHAPE
A transformer is a shape pipeline.
Batched matmul first: one contraction for every (batch, head) pair. Then multi-head attention — eight operations, eight shape contracts, no mystery.
A batched matmul is one matmul per batch element: “bij,bjk->bik” keeps b free and contracts j, applying the same (i, k) product independently to every b. In attention the batch axis is effectively (B, H) — beyond that, every step is a shape you already know: a linear projection, a free reshape, a stride-swapping transpose, a contraction, a softmax over the last axis, and the reverse on the way out.
Concretely: B = 2 sequences, T = 8 tokens, H = 4 heads of D = 16, so the embedding width is E = H · D = 64. The table below is the whole forward pass.
Step
Operation
Shape
Input
X
(B, T, E) = (2, 8, 64)
Project
Q = X @ Wqᵀ, same for K and V
(2, 8, 64)
Split heads
reshape (B, T, H, D) then transpose to (B, H, T, D)
(2, 4, 8, 16)
Scores
einsum “bhtd,bhsd->bhts”, divide by √D
(2, 4, 8, 8)
Weights
softmax over the last axis (each query sums to 1)
(2, 4, 8, 8)
Mix values
einsum “bhts,bhsd->bhtd”
(2, 4, 8, 16)
Merge heads
transpose back, reshape (B, T, E)
(2, 8, 64)
Output
einsum “bte,ek->btk”, concat @ Wo
(2, 8, 64)
The attention shape pipeline
Move the four dials and follow the same nine shapes a transformer computes. Only one tensor is quadratic in the sequence length.
E = H · D = 4 · 16 = 64
X, Q, K, V, context (2, 8, 64)
heads (2, 4, 8, 16)
scores / weights (2, 4, 8, 8)
score entries B·H·T² = 2·4·8² = 512
score cost B·H·T²·D = 8,192 multiply-adds
score memory 2.0 KB (float32)
E = 4·16 = 64, and B·T·E = 1,024
= B·H·T·D = 1,024 — splitting heads regroups the same numbers.
Every step is one of the operations from the chapters: matmul, a free reshape, a stride-swapping transpose, softmax over the last axis, and a contraction. Nothing is magic — just shapes in a row.
Derivation: hand-check one query, and why we divide by √D
Shrink attention to one head, two tokens and D = 2, then compute it by hand. The query is q = [1, 0]; the keys are k₁ = [1, 0] and k₂ = [0, 1]; the values are v₁ = [1, 0] and v₂ = [0, 1].
raw scores = [q·k₁, q·k₂] = [1, 0]
scaled = raw / √D = [1/1.4142, 0] = [0.707, 0]
softmax = [e^0.707, e^0] / (e^0.707 + e^0)
= [2.028, 1.000] / 3.028 = [0.670, 0.330]
context = 0.670 · [1, 0] + 0.330 · [0, 1] = [0.670, 0.330]
shape check = (1, 1, 2, 2) scores → (1, 1, 2, 2) weights → (1, 1, 2, 2) context
if q matched both keys equally, softmax would give [0.5, 0.5] and the
context would be the plain average of the two values.
Why scale by 1/√D? A dot product of D independent entries has spread proportional to √D — with random ±1 entries the standard deviation is exactly √D, so D = 16 gives roughly ±4 and D = 64 gives ±8. Large logits push softmax into a one-hot spike where gradients vanish. Dividing by √D restores a roughly unit scale at every head size, so the same learning rate works for all of them.
head-splitting count check (B=2, T=8, E=64, H=4, D=16):
B·T·E = 2·8·64 = 1024
B·H·T·D = 2·4·8·16 = 1024 same numbers, regrouped ✓
score block: 2·4·8·8 = 512 entries — one 8×8 query-key table per head
Layer
Tensor form
Einsum
Linear layer
Y = X @ Wᵀ + b
“bd,od->bo” + bias
Q/K/V projection
Q = X @ Wq
“bte,ek->btk”
Attention scores
Q Kᵀ / √D
“bhtd,bhsd->bhts”
Attention output
softmax(scores) @ V
“bhts,bhsd->bhtd”
BatchNorm / LayerNorm
(X − μ) / σ · γ + β
element-wise + broadcast
Softmax
exp(x) / Σ exp(x)
element-wise + reduction
08
CHECK YOURSELF
Five questions. Then the terms worth keeping.
Answer before you look. The broadcast and view questions are exactly the ones that decide whether a training script runs or silently learns nothing.
0 / 5 answered · 0 correct
01What does the “shape” of a tensor describe?
02In PyTorch, what layout does an image batch tensor use by default?
03What is the result shape when broadcasting tensors of shape (8, 1, 6, 1) and (7, 1, 5)?
04In the einsum expression “bhtd,bhsd->bhts”, what happens to the index d?
05Why does calling .view() fail on a transposed tensor in PyTorch?
Key terms, demystified
Click a card to swap the lazy description for what it actually means.
Exercises from the lesson
Four problems, including a build-it-yourself einsum and the attention shape tracker. Try first; a worked answer is one click away.
Reshape round-trip: take a tensor of shape (2, 3, 4), reshape it to (6, 4), then (24,), then back to (2, 3, 4). Verify the flat element order is preserved at every step.Show one worked answer
Fill the buffer with range(24) = 0…23 in row-major order. (2, 3, 4) → strides (12, 4, 1); element [1][2][3] sits at flat 1·12 + 2·4 + 3·1 = 23, the last. Reshape to (6, 4): strides (4, 1), same 24 numbers in the same order — [5][3] is still flat 5·4 + 3 = 23. Reshape to (24,): strides (1,), index 23 is flat 23. Reshape back to (2, 3, 4): the order never changed, only the metadata, so every element returns to its original position. Check: the buffer sum is 0+1+…+23 = 276 at every stage.
Broadcasting by hand: add a (3, 1) column to a (1, 4) row and write the (3, 4) result. Then write the rule your broadcast_to(shape) method would follow.Show one worked answer
Right-align (3, 1) with (1, 4): axis 0 is 3 vs 1 → 3; axis 1 is 1 vs 4 → 4. With A = [[1], [2], [3]] and B = [[10, 20, 30, 40]], every output (i, j) reads A[i][0] + B[0][j]: [[11, 21, 31, 41], [12, 22, 32, 42], [13, 23, 33, 43]]. broadcast_to(shape) pads missing leading axes with 1, then for each axis either keeps the size (equal) or virtually repeats the single stored value (size 1); it raises if a non-1 size disagrees. No memory is copied — the output is read through the original strides, with stride 0 on the axes that are repeated.
Build a tiny einsum from scratch: support “i,i->” (dot), “ij,jk->ik” (matmul), “i,j->ij” (outer) and “ij->ji” (transpose). Then check it against numpy.Show one worked answer
Parse the string into input labels and an output label. For each input, the rank is the length of its label. Loop over every assignment of values to the union of all index letters (sizes come from the input shapes). For each assignment, take the product of the indexed elements and add it into the output position built from the output letters. “i,i->” over [1, 2, 3] and [4, 5, 6] visits i = 0, 1, 2 and accumulates 1·4 + 2·5 + 3·6 = 32. “ij,jk->ik” over [[1, 2], [3, 4]] and [[5, 6], [7, 8]] gives [[19, 22], [43, 50]]. “i,j->ij” gives the outer grid, and “ij->ji” just relabels, so [[1, 2, 3], [4, 5, 6]] becomes [[1, 4], [2, 5], [3, 6]]. Every case matches np.einsum because the loop is the definition.
Attention shape tracker: given batch_size B, seq_len T, embed_dim E and num_heads H, print the exact shape at every step of multi-head attention. Verify with B=2, T=8, E=64, H=4.Show one worked answer
D = E / H = 16. Input X: (2, 8, 64) = 1024 numbers. Q, K, V = X @ Wᵀ: (2, 8, 64). Split heads with reshape(B, T, H, D) → (2, 8, 4, 16), then transpose(0, 2, 1, 3) → (2, 4, 8, 16). Scores = Q Kᵀ / √D → (2, 4, 8, 8) = 512 numbers. Weights = softmax(scores, axis=-1): (2, 4, 8, 8), each of the 8 rows per (batch, head) sums to 1. Context = weights @ V → (2, 4, 8, 16) = 1024 numbers. Merge heads with transpose(0, 2, 1, 3).reshape(B, T, E) → (2, 8, 64). Output = concat @ W_o → (2, 8, 64). Check: E = H·D = 4·16 = 64, and 2·8·64 = 1024 = 2·4·8·16 — splitting heads rearranges the same numbers.
Terms this lesson borrows from later lessons (or outside)
You do not need to master these here. Each one gets a proper treatment in its own lesson; the one-line meaning is enough to keep reading. Orange dotted underlines in the prose point back to this list.
attention — The mechanism inside transformers where each token scores every other token with a query · key dot product, then reads from the high scorers. (Transformer lessons)
token — A word or word-piece that a language model reads or writes one at a time. A batch of tokens is what the T axis counts. (Transformer lessons)
normalization — Rescaling inputs to a common range (typically mean 0, spread 1) so no feature dominates. BatchNorm and LayerNorm are reductions plus broadcasts. (Lesson 13)
PyTorch — A deep learning framework: arrays (tensors) with automatic differentiation built in. Its shape semantics are the ones this lesson describes.
clustering — Grouping data points so points in a group are closer to each other than to points outside it. Pairwise-distance tensors make it one argmin. (Lesson 10)
nearest neighbour — Finding the stored point(s) closest to a query under some distance; a k-NN classifier votes among them. (Lesson 14)
pooling — Collapsing an axis by averaging or taking the max — for example turning an image's (H, W) grid into one number per channel. (CNN lessons)
KEEP GOING
A picture is a start. Practice is the rest.
This lesson is a port of an open course. Everything here traces back to it — and the next step is running the code yourself.
Lesson text adapted from AI Engineering from Scratch (Phase 01, Lesson 12) and the Math Foundations Notebook reference build. The animated permutation hero, the shape workbench, 3D broadcast ghost lab, reshape/view memory lab, permute canvas, einsum explorer, attention pipeline and every numeric check are original to this page. Every lab runs in your browser.