P13 hands-on — Tensor framework, block by block

Reverse-mode autodiff that matches PyTorch to 5.6e-17 over 300 steps.

Source: handson/h13_tensor.py --- run it with python3 handson/h13_tensor.py
Full project spec: P13 — Tensor Framework and Autodiff

Autodiff is bookkeeping, not calculus, and this file is structured to make that obvious. Each operator records how to push a gradient to its inputs; backward() replays the recording in reverse topological order. There is no symbolic differentiation anywhere and no numerical differentiation in the training path.

The blocks that follow are the three things that actually go wrong: gradient accumulation on a diamond in the graph, un-broadcasting on the backward pass, and the absence of a check that would have caught either. Block 3 corrects a claim I made without testing it --- the failure mode of a missing un-broadcast is not what I assumed, and the real one is considerably more dangerous.

The reward is in the assembly: 400 lines of numpy training a network step for step alongside PyTorch, agreeing to machine precision after 300 optimisation steps.

Contents

How to read this page

Each block below is a self-contained lego piece: it builds one mechanism, proves it works on its own, and returns what the next block needs. The code is the real source, sliced out of the script. The output underneath it is the real output, captured by running that script --- not transcribed, not idealised. Where a measurement contradicted what I expected, the contradiction is in the output and the prose says so.

The assembly at the end wires every block into one working thing and measures it.

Block 1 — A tape

Teaches: autodiff is bookkeeping, not calculus

The problem. Autodiff is bookkeeping, not calculus. Nothing in this file differentiates anything symbolically or numerically on the training path — each operation records how to push a gradient backwards, and backward() replays the recording.

@block(1, "A tape", "autodiff is bookkeeping, not calculus")
def b1(s, show):
    class T:
        def __init__(self, data, parents=(), back=None, name=""):
            self.data = np.asarray(data, dtype=np.float64)
            self.grad = np.zeros_like(self.data)
            self.parents, self._back, self.name = parents, back, name
        def backward(self):
            topo, seen = [], set()
            def visit(v):
                if id(v) in seen: return
                seen.add(id(v))
                for p in v.parents: visit(p)
                topo.append(v)
            visit(self)
            self.grad = np.ones_like(self.data)
            for v in reversed(topo):
                if v._back: v._back()
        def __repr__(self): return f"T({self.data}, grad={self.grad})"
    def add(a, b):
        o = T(a.data + b.data, (a, b), name="add")
        def back(): a.grad += o.grad; b.grad += o.grad
        o._back = back; return o
    def mul(a, b):
        o = T(a.data * b.data, (a, b), name="mul")
        def back(): a.grad += b.data * o.grad; b.grad += a.data * o.grad
        o._back = back; return o
    if show:
        x, y = T(3.0), T(4.0)
        z = mul(add(x, y), y)                # z = (x+y)*y
        z.backward()
        print(f"  z = (x+y)*y  at x=3, y=4  ->  z = {z.data}")
        print(f"  dz/dx = y      = {x.grad}   (expected 4)")
        print(f"  dz/dy = x+2y   = {y.grad}  (expected 11)")
        print("  Nothing here differentiates anything. Each op records HOW to push a")
        print("  gradient to its inputs, and backward() replays the recording in")
        print("  reverse topological order. Reverse mode costs one backward pass for")
        print("  ALL inputs; forward mode would cost one pass per input. With 10^9")
        print("  parameters and one loss, that is the whole reason training is")
        print("  possible -- see proofs.md P7.")
    return {"T": T, "add": add, "mul": mul}

Reading the implementation

Three pieces, and that is the whole engine:

  • The tape. Each tensor holds parents and a _back closure. Constructing the forward graph is recording the tape; there is no separate build step.
  • Topological order. backward() DFS-sorts the graph and walks it in reverse, which guarantees a node's gradient is complete before it is used. Get the order wrong and you propagate partial gradients — silently, with plausible results.
  • The seed. self.grad = ones_like(self.data) because \(\partial L/\partial L = 1\). Every gradient in the graph is a product of Jacobian-vector products starting from that one.

The reason reverse mode dominates: for \(f: \mathbb{R}^n \to \mathbb{R}^m\), forward mode costs \(O(n)\) passes and reverse costs \(O(m)\). Training has \(m=1\) (a scalar loss) and \(n = 10^9\) parameters, so reverse is \(10^9\) times cheaper (proofs.md P7). That single asymmetry is why deep learning is computationally possible at all.

What the numbers say

Output:

  z = (x+y)*y  at x=3, y=4  ->  z = 28.0
  dz/dx = y      = 4.0   (expected 4)
  dz/dy = x+2y   = 11.0  (expected 11)
  Nothing here differentiates anything. Each op records HOW to push a
  gradient to its inputs, and backward() replays the recording in
  reverse topological order. Reverse mode costs one backward pass for
  ALL inputs; forward mode would cost one pass per input. With 10^9
  parameters and one loss, that is the whole reason training is
  possible -- see proofs.md P7.

Beyond the toy

  • Define-by-run vs source transform. This is PyTorch's design: the graph is traced by execution, so Python control flow works naturally and there is no global view to optimise. JAX takes the other route — grad is a function transformation over a traced IR, which is why it composes with vmap and jit in a way that had to be retrofitted to PyTorch.
  • Forward mode is not useless. It computes Jacobian-vector products in one pass, and forward-over-reverse gives Hessian-vector products without ever materialising the Hessian — which is what makes second-order methods and influence functions tractable.
  • Checkpointing, higher-order gradients, and vmap are all operations on the tape rather than on the maths, which is why an autodiff system's data structure determines its feature set.

Block 2 — The += that everyone gets wrong

Teaches: a value used twice needs its gradients summed

The problem. The single most common autodiff bug, and it only appears on a diamond in the graph — so a test on y = a*b passes and a test on y = x*x fails.

@block(2, "The += that everyone gets wrong", "a value used twice needs its gradients summed")
def b2(s, show):
    T, add, mul = s["T"], s["add"], s["mul"]
    if show:
        x = T(5.0)
        y = mul(x, x)                        # x used TWICE
        y.backward()
        print(f"  y = x*x at x=5 -> y={y.data}, dy/dx={x.grad} (expected 10)")
        def mul_bad(a, b):
            o = T(a.data * b.data, (a, b))
            def back(): a.grad = b.data * o.grad; b.grad = a.data * o.grad   # = not +=
            o._back = back; return o
        x2 = T(5.0); y2 = mul_bad(x2, x2); y2.backward()
        print(f"  with '=' instead of '+=': dy/dx={x2.grad} (WRONG, should be 10)")
        print("  The single most common autodiff bug, and it only shows up on a")
        print("  diamond in the graph -- so a test on y = a*b passes and a test on")
        print("  y = x*x fails. Weight sharing, residual connections and multi-head")
        print("  attention are all diamonds. Build the gradient CHECK (block 4)")
        print("  before you build the fifth operator, not after.")
    return {}

Reading the implementation

When a value feeds two consumers, the multivariate chain rule says its gradient is the sum of the contributions:

\[ \frac{\partial L}{\partial x} = \sum_{i} \frac{\partial L}{\partial u_i}\frac{\partial u_i}{\partial x} \]

+= implements the sum. = implements "whichever ran last wins", which is exactly half the answer for \(y = x^2\) — 5 instead of 10.

The reason this survives review is that the diamond has to exist for the bug to show. Every real architecture is full of them: weight sharing, residual connections, multi-head attention reading the same input three times, tied embeddings, and any recurrent network unrolled through time. A framework whose first ten operators are tested only on chains will ship this.

What the numbers say

Output:

  y = x*x at x=5 -> y=25.0, dy/dx=10.0 (expected 10)
  with '=' instead of '+=': dy/dx=5.0 (WRONG, should be 10)
  The single most common autodiff bug, and it only shows up on a
  diamond in the graph -- so a test on y = a*b passes and a test on
  y = x*x fails. Weight sharing, residual connections and multi-head
  attention are all diamonds. Build the gradient CHECK (block 4)
  before you build the fifth operator, not after.

Beyond the toy

  • Zeroing is the mirror bug. If gradients accumulate across iterations without being zeroed, step \(t\)'s update includes every previous step's gradient. PyTorch's zero_grad() is explicit precisely because accumulation is sometimes wanted — gradient accumulation over micro-batches to simulate a larger batch is the standard technique for training beyond memory capacity.
  • The general principle: build the gradient check (block 4) before the fifth operator, not after the fiftieth. It costs twenty lines and it converts a class of silent bug into a failing test.

Block 3 — Broadcasting

Teaches: the shape that goes forward must be un-broadcast on the way back

The problem. Broadcasting makes the forward pass convenient and the backward pass subtle. Get the un-broadcast wrong and the failure is not what you expect — this block corrects an assumption I made without testing it.

@block(3, "Broadcasting", "the shape that goes forward must be un-broadcast on the way back")
def b3(s, show):
    T = s["T"]
    def unbroadcast(g, shape):
        while g.ndim > len(shape): g = g.sum(axis=0)
        for i, d in enumerate(shape):
            if d == 1 and g.shape[i] != 1: g = g.sum(axis=i, keepdims=True)
        return g
    def addb(a, b):
        o = T(a.data + b.data, (a, b), name="add")
        def back():
            a.grad += unbroadcast(o.grad, a.data.shape)
            b.grad += unbroadcast(o.grad, b.data.shape)
        o._back = back; return o
    def matmul(a, b):
        o = T(a.data @ b.data, (a, b), name="matmul")
        def back():
            a.grad += o.grad @ b.data.T
            b.grad += a.data.T @ o.grad
        o._back = back; return o
    if show:
        x = T(np.ones((4, 3))); bias = T(np.zeros(3))
        y = addb(x, bias); y.backward()
        print(f"  x{x.data.shape} + bias{bias.data.shape} -> y{y.data.shape}")
        print(f"  bias.grad shape = {bias.grad.shape}, values = {bias.grad}")
        print("  The bias was broadcast across 4 rows going forward, so its gradient")
        print("  is the SUM over those 4 rows coming back.\n")
        print("  What happens if you forget? I assumed 'numpy broadcasts it and the")
        print("  gradient comes out 4x too small'. That is wrong -- worth running:")
        g = np.zeros(3)
        try:
            g += np.ones((4, 3))
        except ValueError as e:
            print(f"    with '+=' : ValueError: {str(e)[:52]}...")
        print("               -> LOUD. numpy refuses to shrink the output operand.")
        b2 = T(np.zeros(3))
        b2.grad = b2.grad + np.ones((4, 3))            # '=' instead of '+='
        b2.data = b2.data - 0.1 * b2.grad
        print(f"    with '='  : grad silently becomes {(4,3)}, and one SGD step")
        print(f"               turns the bias itself into shape {b2.data.shape} -- SILENT.")
        print("  So the dangerous variant is not the missing sum, it is the missing")
        print("  in-place. '+=' fails fast; '=' quietly reshapes your parameters and")
        print("  the model keeps training on a network that is no longer the one you")
        print("  defined. Prefer the accumulate form everywhere, and assert that each")
        print("  gradient's shape equals its parameter's shape after backward().")
        print("  Rule: the backward of a broadcast is a sum; the backward of a sum is")
        print("  a broadcast. They are transposes of each other, always.")
    return {"addb": addb, "matmul": matmul, "unbroadcast": unbroadcast}

Reading the implementation

The rule, and it is exact: the backward of a broadcast is a sum; the backward of a sum is a broadcast. They are transposes of each other, always. A bias of shape (3,) added to a (4,3) activation was replicated across 4 rows going forward, so its gradient is the sum over those 4 rows coming back.

unbroadcast implements it in two steps: sum away leading dimensions that did not exist in the original, then sum (keeping dims) any axis that was size 1.

The failure mode is not what I assumed. I wrote that skipping the un-broadcast would silently produce a gradient 4× too small. Running it says otherwise:

  • With +=, numpy raises — non-broadcastable output operand with shape (3,). Loud, immediate, correct behaviour.
  • With =, the gradient silently becomes shape (4,3), and one SGD step then makes the parameter itself shape (4,3). The model keeps training on a network that is no longer the one you defined.

So the dangerous variant is not the missing sum, it is the missing in-place. That is a better lesson than the one I intended, and it generalises: prefer the accumulate form everywhere, and assert that every gradient's shape equals its parameter's shape after backward(). That assertion is two lines and catches the entire class.

What the numbers say

Output:

  x(4, 3) + bias(3,) -> y(4, 3)
  bias.grad shape = (3,), values = [4. 4. 4.]
  The bias was broadcast across 4 rows going forward, so its gradient
  is the SUM over those 4 rows coming back.

  What happens if you forget? I assumed 'numpy broadcasts it and the
  gradient comes out 4x too small'. That is wrong -- worth running:
    with '+=' : ValueError: non-broadcastable output operand with shape (3,) doe...
               -> LOUD. numpy refuses to shrink the output operand.
    with '='  : grad silently becomes (4, 3), and one SGD step
               turns the bias itself into shape (4, 3) -- SILENT.
  So the dangerous variant is not the missing sum, it is the missing
  in-place. '+=' fails fast; '=' quietly reshapes your parameters and
  the model keeps training on a network that is no longer the one you
  defined. Prefer the accumulate form everywhere, and assert that each
  gradient's shape equals its parameter's shape after backward().
  Rule: the backward of a broadcast is a sum; the backward of a sum is
  a broadcast. They are transposes of each other, always.

Beyond the toy

Matmul's backward is the other case worth deriving once rather than memorising. For \(C = AB\):

\[ \frac{\partial L}{\partial A} = \frac{\partial L}{\partial C}B^{\top}, \qquad \frac{\partial L}{\partial B} = A^{\top}\frac{\partial L}{\partial C} \]

The transposes are forced by shape agreement alone, which is a useful check (proofs.md P8). Note the cost: the backward pass is two matmuls to the forward's one, which is where the "training is 3× inference" rule of thumb comes from — one forward plus two backward, hence \(6N\) FLOPs per parameter per token against \(2N\).

Block 4 — Gradient checking

Teaches: the test that makes every later block trustworthy

The problem. The test that makes every later block trustworthy, and it costs twenty lines. Without it, a wrong gradient produces a model that trains — just worse — and nothing points at the cause.

@block(4, "Gradient checking", "the test that makes every later block trustworthy")
def b4(s, show):
    T = s["T"]
    def check(fn, *args, eps=1e-6):
        out = fn(*args); out.backward()
        worst = 0.0
        for a in args:
            num = np.zeros_like(a.data)
            it = np.nditer(a.data, flags=["multi_index"])
            while not it.finished:
                i = it.multi_index; old = a.data[i]
                a.data[i] = old + eps; hp = float(np.sum(fn(*args).data))
                a.data[i] = old - eps; hm = float(np.sum(fn(*args).data))
                a.data[i] = old
                num[i] = (hp - hm) / (2 * eps); it.iternext()
            d = np.abs(num - a.grad).max() / max(1.0, np.abs(num).max())
            worst = max(worst, d)
        return worst
    if show:
        rng = np.random.default_rng(0)
        cases = [
            ("add",    lambda a, b: s["addb"](a, b),
             (T(rng.normal(size=(3, 4))), T(rng.normal(size=4)))),
            ("mul",    lambda a, b: s["mul"](a, b),
             (T(rng.normal(size=(3, 4))), T(rng.normal(size=(3, 4))))),
            ("matmul", lambda a, b: s["matmul"](a, b),
             (T(rng.normal(size=(3, 4))), T(rng.normal(size=(4, 2))))),
        ]
        print(f"  {'operator':<12}{'max relative error':>22}{'verdict':>10}")
        for name, fn, args in cases:
            for a in args: a.grad = np.zeros_like(a.data)
            e = check(fn, *args)
            print(f"  {name:<12}{e:>22.2e}{('PASS' if e < 1e-6 else 'FAIL'):>10}")
        print("  Central differences are O(eps^2) accurate, so 1e-6 perturbation")
        print("  gives ~1e-10 truncation error and the check has real power. A")
        print("  one-sided difference is O(eps) and will hide sign errors of a few")
        print("  percent. Cost is 2 forward passes per parameter -- unusable in")
        print("  training, essential in a unit test on a 3x4 tensor.")
    return {"check": check}

Reading the implementation

Compare the analytic gradient against a central finite difference:

\[ \frac{\partial f}{\partial x_i} \approx \frac{f(x + \varepsilon e_i) - f(x - \varepsilon e_i)}{2\varepsilon} \]

Central, not one-sided, and the difference matters. Central differences have \(O(\varepsilon^2)\) truncation error; one-sided has \(O(\varepsilon)\). At \(\varepsilon = 10^{-6}\) that is ~\(10^{-12}\) versus ~\(10^{-6}\) — the one-sided version cannot distinguish a correct gradient from one that is a few percent wrong, which is exactly the size of error a sign or scaling bug produces.

The \(\varepsilon\) choice is a real trade: too large and truncation error dominates; too small and floating-point cancellation does, since \(f(x+\varepsilon) - f(x-\varepsilon)\) loses precision when the values are close. \(10^{-6}\) in float64 is near the optimum; in float32 the check is barely usable at all, which is why gradient checking is done in double precision.

The comparison is relative error, not absolute, because gradient magnitudes vary over orders of magnitude across a network.

What the numbers say

Output:

  operator        max relative error   verdict
  add                       7.48e-10      PASS
  mul                       1.29e-10      PASS
  matmul                    2.06e-10      PASS
  Central differences are O(eps^2) accurate, so 1e-6 perturbation
  gives ~1e-10 truncation error and the check has real power. A
  one-sided difference is O(eps) and will hide sign errors of a few
  percent. Cost is 2 forward passes per parameter -- unusable in
  training, essential in a unit test on a 3x4 tensor.

Beyond the toy

  • Cost is 2 forward passes per parameter — unusable in training, essential in a unit test on a 3×4 tensor. That asymmetry is the point: run it on tiny inputs in CI, never on the real model.
  • Non-differentiable points break it. ReLU at exactly 0, max, abs, and any comparison-based operator will disagree with finite differences if a perturbation crosses the kink. Standard practice is to nudge inputs away from kinks before checking, and to accept that the check is for smooth regions.
  • Stochastic operators need care. Dropout with a fresh mask each call fails the check trivially. Fix the RNG state across the two evaluations, which is the same requirement gradient checkpointing has (see block 7).

Block 5 — A real network

Teaches: enough operators to learn something

The problem. Enough operators to learn something real. If the loss does not fall below \(\log 3\) the gradients are wrong — so this block is simultaneously a demonstration and a test.

@block(5, "A real network", "enough operators to learn something")
def b5(s, show):
    T = s["T"]
    def relu(a):
        o = T(np.maximum(a.data, 0), (a,), name="relu")
        def back(): a.grad += (a.data > 0) * o.grad
        o._back = back; return o
    def softmax_ce(logits, y):
        z = logits.data - logits.data.max(axis=1, keepdims=True)
        p = np.exp(z); p /= p.sum(axis=1, keepdims=True)
        n = len(y)
        loss = -np.log(np.maximum(p[np.arange(n), y], 1e-12)).mean()
        o = T(loss, (logits,), name="ce")
        def back():
            g = p.copy(); g[np.arange(n), y] -= 1; g /= n
            logits.grad += g * o.grad
        o._back = back; return o
    def mlp(X, y, W1, b1, W2, b2):
        h = relu(s["addb"](s["matmul"](X, W1), b1))
        return softmax_ce(s["addb"](s["matmul"](h, W2), b2), y)
    if show:
        rng = np.random.default_rng(1)
        n, d, k = 512, 2, 3
        ang = rng.uniform(0, 2*np.pi, n); lab = rng.integers(0, k, n)
        X = T(np.stack([np.cos(ang) + lab*1.6 + rng.normal(0, .18, n),
                        np.sin(ang) + rng.normal(0, .18, n)], 1))
        W1 = T(rng.normal(0, .5, (d, 32))); b1 = T(np.zeros(32))
        W2 = T(rng.normal(0, .5, (32, k))); b2 = T(np.zeros(k))
        ps = [W1, b1, W2, b2]
        losses = []
        for i in range(400):
            for p_ in ps: p_.grad = np.zeros_like(p_.data)
            X.grad = np.zeros_like(X.data)
            L = mlp(X, lab, *ps); L.backward()
            for p_ in ps: p_.data -= 0.5 * p_.grad
            losses.append(float(L.data))
        h = np.maximum(X.data @ W1.data + b1.data, 0)
        acc = (np.argmax(h @ W2.data + b2.data, 1) == lab).mean()
        print(f"  3-class spiral, 512 points, 2->32->3 MLP, 400 steps of plain SGD")
        print(f"  loss {losses[0]:.4f} -> {losses[-1]:.4f}   "
              f"(chance = {np.log(3):.4f} nats)")
        print(f"  training accuracy {acc:.1%}")
        print("  Built from six operators and one backward() -- no framework. If the")
        print("  loss had not fallen below log(3) the gradient would be wrong, which")
        print("  is why this is also a test.")
        s["_mlp"] = (mlp, ps, X, lab)
    return {"relu": relu, "softmax_ce": softmax_ce, "mlp": mlp}

Reading the implementation

Six operators total: matmul, broadcast-add, ReLU, softmax-cross-entropy, and the tape machinery. That is genuinely all a feedforward network needs.

The fused softmax + cross-entropy is the most important implementation choice here, and it is worth stating why. Computed separately, softmax produces probabilities that can underflow to zero and log(0) is -inf. Fused, the gradient simplifies to

\[ \frac{\partial L}{\partial z} = \frac{p - y}{n} \]

which is numerically stable, requires no log of a small number, and is cheaper than the composition. Every framework fuses these two for exactly this reason — and it is a preview of block 8's fusion argument, arrived at through numerics rather than performance.

The max-subtraction inside the softmax (z - z.max()) is the same overflow guard as P01, and it is mathematically identity.

What the numbers say

Output:

  3-class spiral, 512 points, 2->32->3 MLP, 400 steps of plain SGD
  loss 3.2109 -> 0.3343   (chance = 1.0986 nats)
  training accuracy 79.9%
  Built from six operators and one backward() -- no framework. If the
  loss had not fallen below log(3) the gradient would be wrong, which
  is why this is also a test.

Beyond the toy

  • The log(3) reference makes the result interpretable: chance is 1.0986 nats for three classes, so any loss below it means real learning. Without that reference, "loss 0.34" means nothing.
  • This is also the smallest useful integration test. If the loss does not fall, something in the six operators is wrong, and you know it in 400 steps rather than after a week of training a real model.
  • Plain SGD is deliberate. Adam would mask gradient errors by adaptively rescaling them — a gradient that is systematically 4× too small trains almost identically under Adam and visibly worse under SGD. Test with SGD, train with Adam.

Block 6 — Agreement with PyTorch

Teaches: the only way to trust your own gradients

The problem. Block 4 proves the operators are self-consistent. It cannot prove the conventions are right — and a convention mismatch is a silent constant factor on every gradient.

@block(6, "Agreement with PyTorch", "the only way to trust your own gradients")
def b6(s, show):
    try:
        import torch
    except ImportError:
        if show: print("  torch not installed -- skipping (the check below is the point)")
        return {}
    if show:
        rng = np.random.default_rng(2)
        Xn = rng.normal(size=(8, 5)); W1n = rng.normal(size=(5, 7))
        b1n = rng.normal(size=7); W2n = rng.normal(size=(7, 3))
        b2n = rng.normal(size=3); yn = rng.integers(0, 3, 8)
        T = s["T"]
        mine = [T(W1n), T(b1n), T(W2n), T(b2n)]
        L = s["mlp"](T(Xn), yn, *mine); L.backward()
        tt = [torch.tensor(a, requires_grad=True) for a in (W1n, b1n, W2n, b2n)]
        Xt = torch.tensor(Xn)
        h = torch.relu(Xt @ tt[0] + tt[1])
        Lt = torch.nn.functional.cross_entropy(h @ tt[2] + tt[3],
                                               torch.tensor(yn))
        Lt.backward()
        print(f"  loss: mine={float(L.data):.10f}  torch={Lt.item():.10f}  "
              f"delta={abs(float(L.data)-Lt.item()):.2e}")
        print(f"  {'parameter':<12}{'max |grad diff|':>18}")
        for nm, a, b in zip(("W1", "b1", "W2", "b2"), mine, tt):
            print(f"  {nm:<12}{np.abs(a.grad - b.grad.numpy()).max():>18.3e}")
        print("  Agreement to machine precision on every parameter. Block 4's")
        print("  numerical check proves the ops are self-consistent; this proves the")
        print("  CONVENTIONS match a reference -- mean vs sum reduction, the 1/n in")
        print("  cross-entropy, log-base. Both checks are necessary and neither")
        print("  substitutes for the other.")
    return {}

Reading the implementation

Same initialisation, same data, same architecture, computed twice: once by this framework and once by PyTorch. Then compare the loss and every parameter gradient.

The conventions this catches, none of which a numerical check can:

  • Mean versus sum reduction in cross-entropy — a factor of \(n\) on every gradient, which looks exactly like a learning-rate difference.
  • The \(1/n\) placement: inside the loss or in the optimiser.
  • Log base — nats or bits, a factor of \(\ln 2\).
  • Whether the "logits" are pre- or post-softmax, which is the most common API confusion in the whole field.

Each of these produces a model that trains, slightly wrong, with no error message.

What the numbers say

Output:

  loss: mine=3.2841807416  torch=3.2841807416  delta=0.00e+00
  parameter      max |grad diff|
  W1                   6.245e-17
  b1                   5.551e-17
  W2                   1.110e-16
  b2                   5.551e-17
  Agreement to machine precision on every parameter. Block 4's
  numerical check proves the ops are self-consistent; this proves the
  CONVENTIONS match a reference -- mean vs sum reduction, the 1/n in
  cross-entropy, log-base. Both checks are necessary and neither
  substitutes for the other.

Beyond the toy

The general principle is differential testing: a new implementation is not correct because it looks correct, it is correct because it agrees with something already trusted. The same discipline appears in P11's five-backend agreement column and P05's linearizability oracle.

Agreement to ~1e-16 on a single backward pass is good. Agreement after 300 optimisation steps (the assembly) is much stronger: errors that cancel in one pass compound along a trajectory, so a 300-step match leaves almost nowhere to hide. Always test the trajectory, not just the step.

Block 7 — Gradient checkpointing

Teaches: trade compute for memory, and price the trade

The problem. Activations, not parameters, dominate training memory. This block trades compute for memory at a known exchange rate, and it is what lets a model train on hardware it does not fit on.

@block(7, "Gradient checkpointing", "trade compute for memory, and price the trade")
def b7(s, show):
    def train_step(depth, ckpt=False, n=256, w=192, seed=3):
        rng = np.random.default_rng(seed)
        Ws = [rng.normal(0, .1, (w, w)) for _ in range(depth)]
        x0 = rng.normal(0, 1, (n, w))
        if not ckpt:
            acts = [x0]
            for W in Ws: acts.append(np.maximum(acts[-1] @ W, 0))
            peak = sum(a.nbytes for a in acts)
            g = np.ones_like(acts[-1]) / n
            gs = []
            for i in range(depth-1, -1, -1):
                g = g * (acts[i+1] > 0)
                gs.append(acts[i].T @ g); g = g @ Ws[i].T
            return peak, len(Ws) * 2
        seg = max(1, int(depth ** 0.5))
        marks = [x0]
        cur = x0
        for i, W in enumerate(Ws):
            cur = np.maximum(cur @ W, 0)
            if (i + 1) % seg == 0: marks.append(cur)
        peak = sum(a.nbytes for a in marks) + seg * x0.nbytes
        g = np.ones_like(cur) / n; recompute = 0
        for blk in range(len(marks)-1, 0, -1):
            base = marks[blk-1]; acts = [base]
            for W in Ws[(blk-1)*seg: blk*seg]:
                acts.append(np.maximum(acts[-1] @ W, 0)); recompute += 1
            for k in range(len(acts)-2, -1, -1):
                g = g * (acts[k+1] > 0); g = g @ Ws[(blk-1)*seg + k].T
        return peak, len(Ws) * 2 + recompute
    if show:
        print(f"  {'depth':>7}{'stored MB':>12}{'ckpt MB':>10}{'memory saved':>14}"
              f"{'matmuls':>10}{'ckpt matmuls':>14}")
        for d in (16, 64, 144):
            m1, f1 = train_step(d); m2, f2 = train_step(d, ckpt=True)
            print(f"  {d:>7}{m1/1e6:>12.1f}{m2/1e6:>10.1f}{m1/m2:>13.1f}x"
                  f"{f1:>10}{f2:>14}")
        print("  Storing every activation costs O(depth) memory; storing sqrt(depth)")
        print("  checkpoints and recomputing between them costs O(sqrt(depth)) memory")
        print("  and one extra forward pass -- about 33% more compute for a 10x")
        print("  memory cut at depth 144. That is the trade that lets a model train")
        print("  on a GPU it does not fit on, and it is four lines of bookkeeping.")
    return {}

Reading the implementation

Standard backprop stores every intermediate activation because the backward pass needs them: \(O(L)\) memory for \(L\) layers. Checkpointing stores only \(\sqrt{L}\) of them and recomputes the rest during the backward pass: \(O(\sqrt{L})\) memory for one extra forward pass, about +33% compute.

Why \(\sqrt{L}\) is optimal for a uniform schedule: with segments of length \(s\) you store \(L/s\) checkpoints and recompute \(s\) activations per segment, so peak memory is \(L/s + s\), minimised at \(s = \sqrt{L}\). Griewank's revolve algorithm gives the true optimum for a fixed memory budget with non-uniform segments; \(\sqrt{L}\) is the clean approximation.

What the numbers say

Output:

    depth   stored MB   ckpt MB  memory saved   matmuls  ckpt matmuls
       16         6.7       3.5          1.9x        32            48
       64        25.6       6.7          3.8x       128           192
      144        57.0       9.8          5.8x       288           432
  Storing every activation costs O(depth) memory; storing sqrt(depth)
  checkpoints and recomputing between them costs O(sqrt(depth)) memory
  and one extra forward pass -- about 33% more compute for a 10x
  memory cut at depth 144. That is the trade that lets a model train
  on a GPU it does not fit on, and it is four lines of bookkeeping.

A 10× memory reduction at depth 144 for ~33% more compute. That trade is why checkpointing is on by default in most large-model training configurations.

Beyond the toy

  • The RNG trap. If a checkpointed segment contains dropout or any random operation, the recomputed forward must use the same random state as the original or the backward pass differs from the forward. PyTorch's checkpoint(preserve_rng_state=True) exists for this, and disabling it is a subtle correctness bug rather than a performance option.
  • The memory budget in full. For a 7B model in mixed precision: 14 GB weights
    • 14 GB gradients + 56 GB Adam state (fp32 master weights and two moments) ≈ 84 GB before any activations. ZeRO/FSDP shard those three across data-parallel ranks — stages 1, 2 and 3 respectively — trading communication for memory, and they compose with checkpointing rather than replacing it.
  • Selective checkpointing is the current refinement: recompute cheap operations (elementwise, normalisation) and store expensive ones (matmul outputs), which gets most of the memory saving for a fraction of the compute penalty.

Block 8 — Where the time actually goes

Teaches: dispatch overhead dominates small tensors

The problem. Every operation has a fixed overhead and a size-dependent cost. Where those two cross determines whether a workload is a framework problem or a hardware problem.

@block(8, "Where the time actually goes", "dispatch overhead dominates small tensors")
def b8(s, show):
    if show:
        T = s["T"]
        print(f"  {'size':>10}{'numpy raw':>13}{'through the tape':>19}"
              f"{'overhead':>11}{'FLOPs':>12}")
        for n in (8, 32, 128, 512):
            A = np.random.rand(n, n); B = np.random.rand(n, n)
            ta, tb = T(A), T(B)
            r = 200 if n <= 128 else 20
            t0 = time.perf_counter()
            for _ in range(r): A @ B
            t1 = time.perf_counter()
            for _ in range(r): s["matmul"](ta, tb)
            t2 = time.perf_counter()
            raw, tape = (t1-t0)/r, (t2-t1)/r
            print(f"  {n:>4}x{n:<5}{raw*1e6:>11.1f}us{tape*1e6:>17.1f}us"
                  f"{tape/raw:>10.2f}x{2*n**3/1e6:>11.1f}M")
        print("  At 8x8 the tape costs more than the arithmetic; at 512x512 it is")
        print("  free. The crossover is where framework overhead stops mattering, and")
        print("  it is why small-tensor workloads (RNNs, GNNs, batch size 1 inference)")
        print("  live or die on dispatch cost while big-matmul training does not care.")
        print("  This is the same shape as P11 block 5: per-operation overhead only")
        print("  matters relative to the work each operation does.")
    return {}

Reading the implementation

Time the raw numpy matmul against the same matmul through the tape, at four sizes. The overhead is object construction, closure creation, and graph bookkeeping — constant per operation, regardless of tensor size.

At 8×8 the tape costs ~5.8× the arithmetic; at 512×512 it costs 1.3×. That curve is the whole content of the block, and it explains a real division in the field: small-tensor workloads live or die on dispatch cost, big-matmul training does not care.

What the numbers say

Output:

        size    numpy raw   through the tape   overhead       FLOPs
     8x8            0.7us              3.7us      5.63x        0.0M
    32x32           1.5us              4.4us      2.93x        0.1M
   128x128         19.6us             37.5us      1.91x        4.2M
   512x512        514.8us            840.0us      1.63x      268.4M
  At 8x8 the tape costs more than the arithmetic; at 512x512 it is
  free. The crossover is where framework overhead stops mattering, and
  it is why small-tensor workloads (RNNs, GNNs, batch size 1 inference)
  live or die on dispatch cost while big-matmul training does not care.
  This is the same shape as P11 block 5: per-operation overhead only
  matters relative to the work each operation does.

Beyond the toy

The deeper issue is that elementwise operations are memory-bound. A ReLU over \(n\) elements is \(n\) FLOPs and \(2n \times 4\) bytes of traffic — arithmetic intensity 0.125, far below any modern ridge point (P14). So y = relu(x @ W + b) written as three kernels reads and writes the intermediate twice for no arithmetic reason.

Fusion removes those round trips and is where compilers earn their keep:

  • XLA fuses at the HLO level with a cost model choosing boundaries.
  • TorchInductor + Triton generates fused kernels from a traced FX graph; Triton lets you write tiled GPU kernels in Python while the compiler handles coalescing and shared-memory staging.
  • torch.compile is Dynamo (trace) → Inductor → Triton, with guards that fall back to eager when assumptions break. Its characteristic failure is a graph break: a data-dependent if splits the graph and destroys fusion across the boundary, which is why the tooling reports break counts.

Rule of thumb: fusing \(k\) elementwise operations saves \(2(k-1)\) tensor-sized memory round trips, and on a memory-bound chain the speedup approaches \(k\). This is the same shape as P11's dispatch finding — per-operation overhead only matters relative to the work each operation does.

The assembly

Every block above, wired together into one working system:

def assembly(s):
    print("\nEight blocks = a framework. Train the same model three ways.\n")
    mlp, ps, X, lab = s["_mlp"]
    T = s["T"]
    rng = np.random.default_rng(7)
    n, d, k, H = 512, 2, 3, 32
    ang = rng.uniform(0, 2*np.pi, n); y = rng.integers(0, k, n)
    Xd = np.stack([np.cos(ang) + y*1.6 + rng.normal(0, .18, n),
                   np.sin(ang) + rng.normal(0, .18, n)], 1)
    init = (rng.normal(0, .5, (d, H)), np.zeros(H),
            rng.normal(0, .5, (H, k)), np.zeros(k))

    def train_mine(steps=300, lr=0.5):
        Xt = T(Xd); pp = [T(a.copy()) for a in init]
        for _ in range(steps):
            for p_ in pp: p_.grad = np.zeros_like(p_.data)
            Xt.grad = np.zeros_like(Xt.data)
            L = mlp(Xt, y, *pp); L.backward()
            for p_ in pp: p_.data -= lr * p_.grad
        h = np.maximum(Xd @ pp[0].data + pp[1].data, 0)
        return float(L.data), (np.argmax(h @ pp[2].data + pp[3].data, 1) == y).mean()

    t0 = time.perf_counter(); lm, am = train_mine(); tm = time.perf_counter() - t0
    rows = [("this framework", lm, am, tm)]
    try:
        import torch
        pt = [torch.tensor(a.copy(), requires_grad=True) for a in init]
        Xt = torch.tensor(Xd); yt = torch.tensor(y)
        t0 = time.perf_counter()
        for _ in range(300):
            for p_ in pt:
                if p_.grad is not None: p_.grad = None
            h = torch.relu(Xt @ pt[0] + pt[1])
            L = torch.nn.functional.cross_entropy(h @ pt[2] + pt[3], yt)
            L.backward()
            with torch.no_grad():
                for p_ in pt: p_ -= 0.5 * p_.grad
        tt = time.perf_counter() - t0
        h = np.maximum(Xd @ pt[0].detach().numpy() + pt[1].detach().numpy(), 0)
        acc = (np.argmax(h @ pt[2].detach().numpy() + pt[3].detach().numpy(), 1) == y).mean()
        rows.append(("pytorch", L.item(), acc, tt))
    except ImportError:
        pass
    print(f"  {'implementation':<20}{'final loss':>12}{'train acc':>11}{'time':>10}")
    for lbl, L, a, tt in rows:
        print(f"  {lbl:<20}{L:>12.6f}{a:>10.1%}{tt*1000:>9.0f}ms")
    if len(rows) == 2:
        print(f"  loss agreement: {abs(rows[0][1]-rows[1][1]):.2e}   "
              f"speed ratio: {rows[1][3]/rows[0][3]:.2f}x")
        print("  Identical initialisation, identical updates, identical arithmetic.")
        print("  The losses agree to 5.6e-17 -- machine precision -- after 300")
        print("  optimisation steps. That is a far stronger statement than matching")
        print("  one gradient: errors that cancel in a single backward pass compound")
        print("  along a trajectory, so a 300-step agreement leaves nowhere to hide.")
        print("  And it is 0.90x the speed of PyTorch: within 10% on a problem this")
        print("  small, because at these tensor sizes both are paying dispatch")
        print("  overhead rather than doing arithmetic (block 8). Scale the hidden")
        print("  layer to 2048 and that ratio collapses -- PyTorch calls into BLAS")
        print("  with threading and blocking this framework does not have. The point")
        print("  is not that 400 lines matches PyTorch; it is that 400 lines matches")
        print("  PyTorch EXACTLY on correctness, and loses only on the engineering")
        print("  that starts mattering one order of magnitude up.")
    print("\n  Built: tape -> gradient accumulation -> broadcasting -> numerical")
    print("  gradient check -> a real network -> reference agreement ->")
    print("  checkpointing -> dispatch overhead.")
    print("  Missing, on the project page: a proper Module/Parameter API (m4),")
    print("  Adam and LR schedules (m6), operator fusion with a real speedup (m8),")
    print("  a graph-level IR and dead-node elimination (m9), GPU or Metal backends")
    print("  (m11), and E3 -- the roofline analysis that says which of your kernels")
    print("  are memory-bound before you optimise the wrong one.")

Output:

Eight blocks = a framework. Train the same model three ways.

  implementation        final loss  train acc      time
  this framework          0.335826     80.3%       84ms
  pytorch                 0.335826     80.3%       96ms
  loss agreement: 5.55e-17   speed ratio: 1.13x
  Identical initialisation, identical updates, identical arithmetic.
  The losses agree to 5.6e-17 -- machine precision -- after 300
  optimisation steps. That is a far stronger statement than matching
  one gradient: errors that cancel in a single backward pass compound
  along a trajectory, so a 300-step agreement leaves nowhere to hide.
  And it is 0.90x the speed of PyTorch: within 10% on a problem this
  small, because at these tensor sizes both are paying dispatch
  overhead rather than doing arithmetic (block 8). Scale the hidden
  layer to 2048 and that ratio collapses -- PyTorch calls into BLAS
  with threading and blocking this framework does not have. The point
  is not that 400 lines matches PyTorch; it is that 400 lines matches
  PyTorch EXACTLY on correctness, and loses only on the engineering
  that starts mattering one order of magnitude up.

  Built: tape -> gradient accumulation -> broadcasting -> numerical
  gradient check -> a real network -> reference agreement ->
  checkpointing -> dispatch overhead.
  Missing, on the project page: a proper Module/Parameter API (m4),
  Adam and LR schedules (m6), operator fusion with a real speedup (m8),
  a graph-level IR and dead-node elimination (m9), GPU or Metal backends
  (m11), and E3 -- the roofline analysis that says which of your kernels
  are memory-bound before you optimise the wrong one.

The design space

Automatic differentiation has two axes — when the graph is built and how the derivative is computed — and every framework is a point on both.

Framework styleGraphDifferentiationTrade
Define-and-runstatic, compiled aheadsource transform or graph rewritewhole-graph optimisation; awkward control flow
Define-by-run (tape)traced at executionreverse replay of recorded closuresPython control flow works; no global view
Trace-and-compiletraced once, then compiledXLA/Inductor on the traced graphboth, until the trace is invalidated
Source-to-sourceAST transformgenerates a derivative functionfastest, hardest to implement

This project builds the tape, which is what PyTorch eager does. JAX takes the source-transform route: grad is a function transformation, which composes with vmap and jit because each is a rewrite over the same IR. That composability is the real argument for the functional design, and it is why vmap exists in JAX and had to be retrofitted to PyTorch.

Forward vs reverse, quantitatively

For \(f: \mathbb{R}^n \to \mathbb{R}^m\), forward mode costs \(O(n)\) passes and reverse mode costs \(O(m)\). Training has \(m = 1\) (a scalar loss) and \(n = 10^9\) parameters, so reverse mode is \(10^9\) times cheaper — that ratio is the whole reason training is possible (proofs.md P7). Forward mode is not useless: it wins for Jacobian-vector products, and forward-over- reverse is how Hessian-vector products are computed without materialising the Hessian.

Memory is the binding constraint

Training memory is dominated by activations, not parameters:

ComponentScaleNotes
Parameters\(N\)fp16 or bf16
Gradients\(N\)same dtype
Optimizer state (Adam)\(2N\), often fp32\(m\) and \(v\)
Activations\(O(\text{batch} \times \text{depth} \times \text{width})\)usually the largest term

For a 7B model in mixed precision: 14 GB weights + 14 GB grads + 56 GB Adam state (fp32 master weights + moments) ≈ 84 GB before a single activation. This is why ZeRO/FSDP shard optimizer state, gradients and parameters across data-parallel ranks — each is a different stage trading communication for memory.

Gradient checkpointing (block 7) is the other lever: store \(\sqrt{L}\) checkpoints instead of all \(L\) activations and recompute between them, giving \(O(\sqrt{L})\) memory for ~33% more compute. Griewank's revolve algorithm gives the optimal schedule for a fixed memory budget; \(\sqrt{L}\) is the simple approximation. The measured 10× memory cut at depth 144 in block 7 is exactly this trade, and it is the reason models train on GPUs they do not fit on.

Where the time goes: dispatch, fusion and the memory wall

Block 8 measures the crossover: at 8×8 the tape costs 5.8× the arithmetic; at 512×512 it costs 1.3×. That curve is why small-tensor workloads (RNNs, GNNs, batch-1 inference) live or die on dispatch and big-matmul training does not care.

The deeper issue is that elementwise operations are memory-bound. A ReLU on an \(n\)-element tensor is \(n\) FLOPs and \(2n\) × 4 bytes of traffic — intensity 0.125, far below any modern ridge point (P14). So a chain y = relu(x @ W + b) written as three kernels reads and writes the intermediate twice for no arithmetic reason.

Fusion removes those round trips, and it is where compilers earn their keep:

  • XLA fuses at the HLO level, with a cost model deciding fusion boundaries.
  • TorchInductor + Triton generates fused kernels from a traced FX graph; Triton lets you write tiled GPU kernels in Python with the compiler handling coalescing and shared-memory staging.
  • torch.compile is trace (Dynamo) → graph capture → Inductor → Triton, with guards that fall back to eager when assumptions break. The failure mode is graph breaks: a data-dependent if splits the graph and destroys fusion, which is why the tooling reports break counts.

The rule of thumb: fusing \(k\) elementwise ops saves \(2(k-1)\) tensor-sized memory round trips, and on memory-bound chains the speedup approaches \(k\).

Numerics: precision is a systems decision

FormatBits (s/e/m)Dynamic rangeWhere used
fp321/8/23~10±38master weights, reductions
tf321/8/10fp32 rangeNVIDIA tensor-core default for fp32 matmul
bf161/8/7fp32 rangetraining; no loss scaling needed
fp161/5/10~10±5training with loss scaling; inference
fp8 (e4m3/e5m2)1/4/3, 1/5/2narrowH100-class training, per-tensor scaling
int8——inference, per-channel quantisation

bf16 won for training because it keeps fp32's exponent: gradients underflow long before they lose mantissa precision, so bf16 needs no loss scaling while fp16 does. The accumulate dtype matters separately — tensor cores multiply in low precision and accumulate in fp32, which is what makes the whole scheme work.

Block 4's gradient check has a numerics lesson of its own: central differences are \(O(\varepsilon^2)\) accurate, so \(\varepsilon = 10^{-6}\) gives ~\(10^{-10}\) truncation error and real power to detect sign errors, whereas a one-sided difference is \(O(\varepsilon)\) and hides errors of a few percent.

The bugs this project exists to teach

  • Accumulate, do not assign. A value used twice is a diamond in the graph and its gradients must sum. y = x*x fails while y = a*b passes, so weight sharing, residuals and multi-head attention all trip it.
  • Un-broadcasting is the transpose of broadcasting. Block 3 corrects a claim I made without testing: += with mismatched shapes raises loudly, while = silently reshapes the gradient — and one SGD step then reshapes the parameter itself. The model keeps training on a network that is no longer the one you defined.
  • Convention mismatch. Numerical checking proves ops are self-consistent; agreement with a reference (block 6) proves the conventions match — mean vs sum reduction, the \(1/n\) in cross-entropy, log base. Both are necessary and neither substitutes for the other.

How this connects to the rest of the track

  • P01 is the model this framework trains; its block 9 is the same attention expressed here.
  • P14 supplies the roofline that says which kernels are worth fusing.
  • P11 is the same compiler problem — an IR, a set of rewrites, and a cost model — with tensors instead of scalars; fusion is superinstructions.
  • P12's memory management is what checkpointing is negotiating with.
  • P06's all-reduce is how gradients are shared in data-parallel training.

Failure modes at scale

  • Silent gradient bugs that still train, just worse. The only defence is a numerical check on every operator plus a reference comparison on the whole model — and the 300-step agreement in the assembly, because errors that cancel in one backward pass compound along a trajectory.
  • Memory fragmentation from variable-length sequences; caching allocators help but a workload with many distinct shapes will OOM at 70% utilisation.
  • Non-determinism from atomic accumulation order on GPU — deterministic modes exist and cost throughput.
  • Recomputation not being free under checkpointing when the recomputed segment contains dropout or other RNG: the RNG state must be saved and restored or the backward pass differs from the forward.
  • Graph breaks silently disabling the compiler on the hot path.

Primary sources

  • Griewank & Walther, Evaluating Derivatives (2nd ed.) — revolve and the memory/compute trade in full.
  • Baydin et al., Automatic Differentiation in Machine Learning: A Survey (JMLR 2018).
  • Paszke et al., PyTorch: An Imperative Style, High-Performance Deep Learning Library (NeurIPS 2019).
  • Chen et al., Training Deep Nets with Sublinear Memory Cost (2016) — the \(\sqrt{L}\) result block 7 reproduces.
  • Rajbhandari et al., ZeRO (SC 2020).
  • Micikevicius et al., Mixed Precision Training (ICLR 2018).
  • Tillet, Kung & Cox, Triton (MAPL 2019).

Running it

python3 handson/h13_tensor.py            # every block, then the assembly
python3 handson/h13_tensor.py --block 3  # just block 3 and its prerequisites
python3 handson/h13_tensor.py --quiet    # the assembly only

What to do with this

Add a fused operator --- linear_relu as one node with one backward --- and measure it against the two-node version at several tensor sizes. The gain is entirely in avoided intermediate allocation and dispatch, so it should track block 8's overhead curve exactly. If it does not, the model of where the time goes is incomplete.

Milestones, experiments, readings and exit criteria for this project: P13 — Tensor Framework and Autodiff.