A "tensor" is just a NumPy array with extra dimensions — and yet 90% of ML bugs are shape mismatches. Master (B, T, D), learn einsum once, and you'll read any paper or any PyTorch repo without squinting. Tensors are how GPUs talk; einsum is the dictionary.
Learning Objectives
After this lesson, you will be able to:
Read and write tensor shapes like (B, T, D) and (B, C, H, W) and explain what each axis means in real ML systems
Translate any matrix or tensor operation into Einstein-summation (einsum) string syntax — and back again
Apply broadcasting rules to combine tensors of different shapes without writing explicit loops
Explain a computational graph as a directed acyclic graph of operations, and describe why reverse-mode autodiff is the algorithm of choice for training neural networks
How loss.backward() does its magic in one line — PyTorch is secretly building a graph of every operation you do during the forward pass. When you call .backward(), it walks that graph in reverse and applies the chain rule node by node. This lesson makes that graph visible
Build this --> Build a "Einsum Decoder" that takes any einsum string ("bhid,bhjd->bhij" or "ii->" or "bik,bkj->bij") and prints a plain-English explanation of what it computes, plus a concrete numerical example with small shapes so you can verify the result by hand
A tensor is just a list of lists of lists. That is genuinely all it is. The reason it gets a fancy name is that ML code constantly stacks more "list of" layers — a single image is a 2D list, a batch of images is a 3D list, a batch of color images is a 4D list, and a batch of color video frames is a 5D list. The math stays the same; only the bookkeeping gets harder.
The shape is the most important thing about a tensor. When you read or write ML code, you are constantly tracking shapes in your head: "this layer takes (B, T, D) and returns (B, T, D)," "this loss reduces (B,) to a scalar," "this projection takes (B, T, D_in) to (B, T, D_out)." Half of all DL bugs are shape mismatches. Get the shapes right and the math usually follows.
(B, D) — a batch of B feature vectors of dimension D. Tabular data, classifier logits.
(B, T, D) — transformer activations. B sequences, T tokens, D-dimensional embeddings. Vary T per batch and you get padding masks.
(B, C, H, W) — CNN inputs/activations in PyTorch (channels-first). B=32, C=3, H=224, W=224 is the canonical ImageNet shape.
(B, H, W, C) — same, channels-last (TensorFlow default, also faster on Apple Silicon).
(B, H, T, D_h) — multi-head attention activations. H is the number of heads, D_h = D / H is per-head dimension.
(L, B, D) — RNN/LSTM convention: time-major instead of batch-major.
The shape tells you nothing about meaning by itself. (32, 3, 224, 224) could be a batch of color photos OR 32 separate 3-channel volumetric scans of size (224, 224). Convention and context decide.
Try it! Open the Python REPL (bottom-right of the screen: click Quick Actions, then Python) and type these lines yourself.
What Do You Think?
A multi-head attention layer for a batch of B=8 sequences of length T=64 with hidden size D=512 and H=8 heads splits D into per-head dimension d_h = D/H = 64. What is the shape of the queries tensor Q inside the attention computation?
Look at that formula carefully. There is a deep symmetry hiding in the indices: i and j appear once on the right-hand side and survive to the left-hand side. k appears twice on the right and not at all on the left — that is precisely the index we summed over. This is the rule that makes Einstein summation possible: any index that appears twice on the right and is missing from the left gets implicitly summed.
Einstein invented this notation in 1916 because writing Σ symbols all day in general relativity equations was driving him crazy. The convention: drop the Σ and let the repeated index do the work.
Cij=AikBkj(Einstein summation: sum over k implied)
The einsum string 'ik,kj->ij' is just the same equation typed sideways. Read it left-to-right:
'ik' — A has axes labeled i (rows) and k (cols).
',' — separator between operands.
'kj' — B has axes labeled k (rows) and j (cols).
'->' — separator before the output spec.
'ij' — the output has axes i and j.
Now apply the einsum semantics to figure out what to do:
Any index appearing on both input sides is a contraction axis (will be multiplied then summed). Here: k.
Any index that appears in inputs but not output also gets summed.
Any index that appears in inputs and output is preserved (free axis).
By those rules, 'ik,kj->ij' means: pair up A and B by k, multiply, sum over k, keep i and j distinct. That is matrix multiplication.
Toggle through the presets below to see einsum operations execute one cell at a time. Watch the contraction index get summed away, and the free indices fall into their slots in the output tensor. Internalise the visual once and you will never have to reason about einsum from scratch again.
In multi-head attention, queries Q and keys K both have shape (B, H, T, d_h) — batch, heads, sequence positions, per-head dimension. The attention-score matrix is shape (B, H, T, T): how much each query position attends to each key position, separately for each head and each batch.
Sbhij=dh1d∑QbhidKbhjd
You will see this exact einsum string in every transformer implementation on GitHub. Memorize it. The next operation — applying the attention weights to the values V of shape (B, H, T, d_h) — is 'bhij,bhjd->bhid': the attention matrix (T, T) weights (T, d_h) to produce (T, d_h) per head per example, contracting over the sequence index j.
Open the visualiser one more time and step through the "Attention QKᵀ" preset (or the closest match). Watch the per-head dimension d collapse while batch b, head h, query position i, and key position j survive into the output. Every "scaled dot-product attention" diagram you have ever seen is this one contraction with a 1/√d_h scaling and a softmax tacked on the end.
Loading visualization...
Quick check
A tensor A has shape (32, 16, 8, 64) and B has shape (32, 16, 64, 8). What is the output shape of einsum('bhij,bhjk->bhik', A, B)?
Does the order of indices on the right-hand side of an einsum string matter? Compare einsum('ik,kj->ij', A, B) vs einsum('ik,kj->ji', A, B).
A whole einsum cookbook in one runnable cell. Each pattern below shows you the einsum string, what it computes, and a shape sanity-check. After running, edit the strings -- swap letters, drop indices, see what breaks.
Sometimes you want to combine tensors of different shapes — say, add a per-feature bias b of shape (D,) to every row of an activation matrix X of shape (B, D). Writing this as a Python loop would be excruciating. Numpy and PyTorch handle it via broadcasting rules.
The core rule: aligned-from-the-right, two shapes are compatible if every paired axis is either equal, or one of them is 1. A "1" axis gets virtually expanded (no memory copy) to match the other tensor.
You have token embeddings X of shape (B, T, D) and positional encodings P of shape (T, D) — same for every example in the batch. Adding them must "broadcast" P across the batch axis.
A common pattern: you have a query q of shape (B, D) and a memory M of shape (T, D) and you want pairwise distances ‖q - M_t‖ for every batch element and every memory slot. You need an output of shape (B, T, D) (per-coordinate differences) before reducing to (B, T).
You have logits L of shape (B, V) (batch B, vocabulary V) and a temperature t of shape (B,) -- a different temperature per example. You write softmax(L / t). What goes wrong if you do not reshape t first?
Once you start composing many tensor operations — Y = (A @ X + b).relu(); Z = Y @ W; loss = (Z - target).pow(2).mean() — there is a hidden structure: a directed acyclic graph (DAG) where nodes are tensors and edges are the operations that produce them.
The forward pass is a topological evaluation of this graph: visit nodes in dependency order, compute each from its parents. The result of the last node is your output (e.g., loss).
The backward pass is the same graph traversed in reverse. Starting from loss (with seed gradient dloss/dloss = 1), each node uses its own local derivative rule plus its already-computed children's gradients to compute the gradient for its parents. Apply the chain rule node-by-node, and every leaf node (your model parameters) ends up with a populated gradient.
∂θi∂L=paths p from θi to L∑e∈p∏∂in(e)∂out(e)
Both modes implement the chain rule. The difference is direction and cost.
Forward mode computes derivatives in the same direction as the forward pass: it pushes a tangent vector forward through the graph. Cost: ~one forward pass per input variable. Cheap if you have few inputs and many outputs.
Reverse mode (a.k.a. backpropagation) traverses the graph backward from a single output: it pulls a cotangent vector back through the graph. Cost: ~one forward pass total, regardless of how many inputs there are.
Neural networks have millions to billions of input parameters and one scalar output (the loss). That makes reverse mode a no-brainer: compute all gradients in roughly the cost of a single forward pass, instead of millions of forward passes.
What loss.backward() actually does
PyTorch's loss.backward() walks the computational graph in reverse, applying the chain rule node by node — this lesson's DAG is exactly the graph being walked. Every tensor with requires_grad=True has its .grad field populated with the corresponding partial derivative. The optimizer (SGD, Adam, etc.) then reads .grad and updates the parameter.
This is also why nn.Embedding(vocab, dim, sparse=True) exists: the gradient of the loss w.r.t. an embedding row is sparse — only the rows for tokens that appeared in the batch get nonzero gradients (the others were never read during the forward pass, so the chain rule yields zero). A naive dense gradient tensor would be (vocab_size, dim) — for a 50K-token vocabulary that's a lot of zeros to allocate. The sparse version stores only the rows that changed, dramatically cutting memory and compute for large embedding tables.
Tensor = list of lists of lists. Rank counts how many "list of" layers; shape is the size at each axis. Almost every ML bug is a shape mismatch — train your eye to read shapes.
Einsum compresses the chain rule notation. Repeated index → sum. Missing-from-output → sum. Distinct → preserved. Three rules, infinite operations.
The attention score 'bhid,bhjd->bhij' is not jargon. It's the chain rule of matrix multiplication, written in compact notation. Every transformer paper uses it.
Broadcasting eliminates loops. Right-aligned shape compatibility lets you add a (D,) bias to a (B, D) batch, or compare a (B, 1, D) query to a (1, T, D) memory, without copying tensors. But it can hide bugs — assert shapes before losses.
Computational graphs make autodiff possible. Forward = topological evaluation of a DAG; backward = reverse-mode chain rule on the same DAG. PyTorch builds it dynamically; JAX traces it once. loss.backward() is just one DAG traversal.
An einsum string is `'bhid,bhjd->bhij'`. Which axis is contracted?
Next up: Eigenvalues and SVD — having mastered the bookkeeping of tensors, we now look at what makes some directions "special" inside a matrix transformation. The eigenvectors of a covariance matrix are the principal components of your data; the singular values reveal the rank structure that LoRA and PCA exploit.