AI01: Autograd from scratch: backpropagation is the chain rule

Prerequisites and observable goals#

Level: AI Lab, Module 1 of 6. About 18 hours over two weeks (eight sessions). Python 3.14 standard library only, so there is nothing to install.

Warm-up: spaced review (10 minutes, closed book)#

If your error log has entries due for review, use two or three of those. Otherwise use these three, drawn from the friction points in LM, the slips learners most often make in Calculus I. All three come back in this module.

1. Differentiate $16x^{-3}$. (Bookkeeping: the dropped negative exponent.)
2. Differentiate $(3x^2+1)^5$. (Nesting dolls.)
3. Compute $\tfrac15-1$, then differentiate $x^{1/5}$. (Fractions under pressure.)

The answers are at the very end of this chapter’s solutions. Number 1 returns in the power operation and in division (lab step 5). Number 2 is backpropagation itself (section 2.3). Number 3 returns in fractional powers (lab step 5 tests $x^{0.5}$).

1. Quick summary#

The running example $L=(a\cdot b+b)\cdot c$, drawn as a computation graph. Black numbers flow forward (inputs to output). Orange numbers flow backward: each is $\partial L/\partial(\text{that value})$, how much $L$ moves per unit nudge of that value. Because $b$ is used twice, it collects 15 along one path and 5 along the other: 20.

Five ideas carry this module.

1. Training asks one question per knob. A neural network is a big expression with adjustable numbers called parameters. Training asks, for each one: if I turn this knob a little, does the error go up or down, and how fast? That is a derivative, one per knob.
2. A partial derivative freezes the other inputs. With many inputs, take the ordinary Calculus I derivative in one input while treating every other input as a constant.
3. Backpropagation is the chain rule, applied node by node. Each operation knows only its own local derivative. It multiplies the gradient arriving from the output side by that local derivative and passes the result to its inputs. A value used in two places adds what arrives from both.
4. Order matters. A node may pass its gradient on only after every node that uses it has reported. That order is a reverse topological sort.
5. Gradient descent turns gradients into learning. Nudge every parameter a small step against its gradient, then repeat. A finite-difference check proves the gradients are right.

What you will build. About 150 lines of Python in the companion code’s ai-lab/m01_autograd folder: engine.py (the Value class and its backward pass), gradcheck.py, nn.py (neurons, layers, networks), and train.py. Seventy tests check them step by step. At the end, a network you wrote learns XOR, a function no single straight line can separate. A companion-code tool, ai-lab/tools/viz.py, animates your own engine’s backward pass, and the Gradient Explorer on the companion website (bezhig.com/book/explorers/ai01-gradient-explorer.html) lets you drag the learning rate, the finite-difference step, and a neuron’s input and watch what happens.

2. Deep explanation#

2.1. The question autograd answers#

Here is the smallest network worth the name: one neuron. It takes two inputs $x_1$ and $x_2$, multiplies each by a weight, adds a bias, and squashes the result:

One neuron. The weights $w_1,w_2$ and the bias $b$ are the knobs. The function $\tanh$ (blue) squashes any number into the range $-1$ to $1$. Its slope (dashed orange) is 1 at $z=0$ and nearly 0 once $z$ is far from 0, which is where learning stalls (AI01.09).

The function $\tanh$ (the hyperbolic tangent) is defined by $\tanh z=\dfrac{e^z-e^{-z}}{e^z+e^{-z}}$. It is an S-shaped curve with $\tanh 0=0$ that levels off at $1$ for large $z$ and at $-1$ for very negative $z$. You will derive its derivative in AI01.09.

The inputs come from data. The weights $w_1$ and $w_2$ and the bias $b$ are the knobs, called parameters, and training chooses them. To judge a choice we need one number that says how wrong the network is, the loss. For one example with target $y$, a common loss is the squared error $L=(\text{output}-y)^2$. Training means changing $w_1$, $w_2$, and $b$ to make $L$ smaller.

For each knob the useful question is: if I nudge this knob up a little, how much does $L$ change? That is the derivative of $L$ with respect to that knob. A real model has millions or billions of knobs. Autograd (automatic gradient) computes every one of those derivatives with one forward pass and one backward pass. That is the machine you build in this module.

2.2. Partial derivatives: freeze the others#

In Calculus I every function had one input. A loss has many. The fix is almost disappointingly simple: to see how $f$ changes with $x$, hold every other input fixed. Then $f$ is an ordinary one-input function of $x$, and you take its ordinary derivative.

Worked example. Let $f(x,y)=x^2y+3y$ at the point $(1,3)$. First, $f(1,3)=1\cdot3+9=12$.

• $\partial f/\partial x$: freeze $y$. Then $y$ is a constant multiplier on $x^2$, and $3y$ is a constant term whose derivative is $0$:

  $$\frac{\partial f}{\partial x}=y\cdot\frac{d}{dx}\bigl(x^2\bigr)+0=2xy,\qquad\text{at }(1,3):\ 2(1)(3)=6.$$

• $\partial f/\partial y$: freeze $x$. Then $x^2$ is a constant coefficient on $y$:

  $$\frac{\partial f}{\partial y}=x^2\cdot1+3=x^2+3,\qquad\text{at }(1,3):\ 1+3=4.$$

Left: the surface $f(x,y)=x^2y+3y$ over the $xy$-plane. Freezing $y=3$ slices it along the blue curve; freezing $x=1$ slices it along the orange line $f(1,y)=4y$. Right: the blue slice is an ordinary Calculus I graph, $g(x)=3x^2+9$, whose tangent at $x=1$ has slope 6. A partial derivative is the slope of a slice.

The gradient. Collect all the partial derivatives into one list: $\nabla f=\left(\dfrac{\partial f}{\partial x},\dfrac{\partial f}{\partial y}\right)$, read “grad $f$.” At $(1,3)$ it is $(6,4)$. For training we use it one entry at a time: each entry tells its own knob which way to turn and how strongly. (Calculus III proves that the whole list, read as an arrow, points in the direction of steepest increase.)

Small nudges add up. For small nudges $\Delta x$ and $\Delta y$,

Each partial predicts the change caused by its own nudge, and the two predictions add. This is the tangent-line approximation of Module 13, with one slope per input. Check it: move from $(1,3)$ to $(1.01,\,3.02)$. The prediction is $6(0.01)+4(0.02)=0.06+0.08=0.14$, so $f\approx12.14$. Exactly, $f(1.01,\,3.02)=(1.0201)(3.02)+9.06=3.080702+9.06=12.140702$. The error, $0.0007$, is far smaller than either nudge. This “add the effects of small nudges” rule is the whole reason backpropagation works (section 2.4).

Check yourself: for the same $f$, what is $\partial f/\partial x$ at $(2,5)$? (Answer: $2\cdot2\cdot5=20$.)

A common misreading. “Freeze $y$” does not mean $y$ stops mattering. The partial derivative $\partial f/\partial x=2xy$ still contains $y$: at $y=3$ the slope in $x$ is 6, and at $y=5$ it would be 10. Freezing is about which input you nudge, not which inputs the answer may depend on. That is why backprop needs every value from the forward pass.

Say it back (one sentence, out loud or on paper): what does $\partial f/\partial y=4$ at $(1,3)$ promise about $f$ when $y$ moves from 3 to 3.01?

2.3. The chain rule on a graph#

Computation graphs. Any expression can be drawn as a graph. Each node is one value. An arrow from $u$ to $v$ means “$u$ was used to compute $v$.” The inputs are the leaves, and the last node is the output. The running example $L=(a\cdot b+b)\cdot c$ with $a=3$, $b=-2$, $c=5$ builds three nodes (the figure in section 1): $e=a\cdot b=-6$, then $d=e+b=-8$, then $L=d\cdot c=-40$. These are the Russian nesting dolls of composition laid out flat, one node per doll.

Local derivatives: the Modular Workbench. Each node knows only how its own output moves when one of its inputs moves, with its other inputs frozen. These are partial derivatives of a tiny function, and there are only a handful of kinds:

Treat these like $u$, $v$, $u'$, $v'$ in the product rule: find each local derivative on its own line, check it, and only then assemble.

The chain rule, graph form. In Calculus I, if $L$ depends on $u$ and $u$ depends on $x$, then $\dfrac{dL}{dx}=\dfrac{dL}{du}\cdot\dfrac{du}{dx}$. On a graph this becomes one rule applied at every node:

The upstream gradient is how much the final output cares about this node; your code stores it in the node’s grad. The backward pass starts at the output with $\partial L/\partial L=1$: nudging $L$ changes $L$ by exactly the nudge. That 1 is called the seed.

Worked backward pass for the running example. Work from the output back toward the leaves.

So $\partial L/\partial a=-10$ and $\partial L/\partial c=-8$. The leaf $b$ received two messages, 5 and 15; the next section shows why they add to $\partial L/\partial b=20$. Check algebraically: $L=(ab+b)c=bc(a+1)$, so $\partial L/\partial a=bc=-10$, $\partial L/\partial b=c(a+1)=20$, and $\partial L/\partial c=b(a+1)=-8$.

Your turn: a faded example. Take $L=(x+y)\cdot x$ at $x=2$, $y=3$, with one inner node $s=x+y=5$, so $L=s\cdot x=10$. Fill in the blanks on paper before reading on.

Then total each leaf. Check after you finish: $s$ gets 2 and $x$ gets 5; the add node passes 2 to both. So $\partial L/\partial x=5+2=7$ and $\partial L/\partial y=2$. Algebra agrees: $L=x^2+xy$ gives $2x+y=7$ and $x=2$. Notice that $x$ is used twice, which is the subject of the next section.

Three patterns cover most of backpropagation. Here $g$ is the gradient arriving at the node from the output side.

2.4. Used twice? Add the gradients#

Why should $b$’s two messages add? The nudge rule from section 2.2 answers it. The node $d=e+b$ depends on $b$ twice: directly, and through $e=a\cdot b$. Nudge $b$ by a small amount $\Delta b$ and follow both effects:

• $e$ changes by $\Delta e\approx\dfrac{\partial e}{\partial b}\Delta b=a\,\Delta b=3\,\Delta b$;
• $d=e+b$ changes by both: $\Delta d=\Delta e+\Delta b\approx3\,\Delta b+\Delta b=4\,\Delta b$;
• $L$ changes by $\Delta L\approx\dfrac{\partial L}{\partial d}\Delta d=5\cdot4\,\Delta b=20\,\Delta b$.

So $\partial L/\partial b=20=15+5$: one term for each path from $b$ to $L$. Check with numbers: $b=-2\to-1.999$ gives $L=\bigl(3(-1.999)+(-1.999)\bigr)(5)=(-7.996)(5)=-39.98$, a change of $0.02=20\times0.001$.

The chain rule as a gearbox. Each local derivative is a gear ratio: turn $b$ by one notch and $e$ turns 3 notches, because $\partial e/\partial b=a=3$. Along a path the ratios multiply. Where two paths meet at an add node, their turns add, the way the differentials in old mechanical computers added two shafts’ rotations.

In general, $\partial L/\partial x$ is the sum, over every path from $x$ to $L$, of the product of the local derivatives along that path. Backpropagation computes exactly that sum without listing the paths. That matters: a graph can have exponentially many paths (a chain of $k$ diamonds like the one in lab step 3 has $2^k$), yet backprop touches each edge once, because each node stores its total so far and hands it on. Reusing stored partial totals instead of recomputing them is called dynamic programming, a technique Module 3 makes central.

Say it back: using the gearbox picture, explain in two sentences why $\partial L/\partial b$ is 20 and not 15.

The smallest example is $y=a\cdot a$ at $a=3$. The multiply node sees the same value on both inputs. It sends “the other input times $g$” to its first input and to its second input, and both are $a$: $3+3=6=\frac{d}{da}a^2$ at $a=3$. With = you would get 3. Lab step 4 has you break your own code this way and watch it happen.

2.5. The order of work: topological sort#

A node’s gradient is finished only when every node that uses it has sent its message. In the running example, $b$’s gradient is not complete until both $d$ and $e$ have reported. So a node may run its backward step only after all of its users have run theirs.

An ordering in which every node comes after all of its inputs is called a topological order. The forward pass already happens in one. The backward pass runs in the reverse of one: output first, and each node before its inputs.

The algorithm: depth-first search, post-order.

1. To visit a node $v$: if $v$ has not been visited, mark it visited, visit each of its inputs, and then append $v$ to the list.
2. Visit the output.

Because $v$ is appended only after all of its inputs, every node lands after its inputs. Reverse the list and every node comes before its inputs.

Trace on the running example (children visited in the order $d,c$ for $L$; $e,b$ for $d$; $a,b$ for $e$): visit $L$ $\to$ visit $d$ $\to$ visit $e$ $\to$ visit $a$ (append $a$) $\to$ visit $b$ (append $b$) $\to$ append $e$ $\to$ visit $b$ again (already visited, skip) $\to$ append $d$ $\to$ visit $c$ (append $c$) $\to$ append $L$.

The post-order list for the running example, and the direction the backward pass walks it.

Say it back: why may $e$ not run its backward step before $d$ has run its own?

The visited set is what stops $b$ from being appended twice. That would be harmless for a leaf, but a shared interior node listed twice would run its backward step twice and double-count its gradient (bug B5).

Cost. Every node and every edge is touched once, so the sort is $O(V+E)$ for $V$ nodes and $E$ edges: proportional to the size of the graph. The whole backward pass costs a small constant multiple of the forward pass, however many inputs there are. (In an engine like yours it is even cheaper than the forward pass: for a 100,000-operation chain on one desktop PC, a complete engine of this design took about 0.06 s for the backward pass, while building the graph took 0.08 to 0.2 s.) That is why training is affordable at all.

A Python limit. Python stops a chain of nested function calls at about 1000 by default (python -c "import sys; print(sys.getrecursionlimit())" prints 1000 on a standard Python installation). A recursive visit on a 5000-node chain therefore raises RecursionError. Stretch S1 replaces the recursion with an explicit stack.

2.6. Gradient descent: turning gradients into learning#

Now use the gradients. Take a knob $w$ with gradient $g=\partial L/\partial w$. If $g>0$, increasing $w$ increases $L$, so decrease $w$. If $g<0$, increasing $w$ decreases $L$, so increase $w$. Both cases say the same thing: step against the gradient.

The Greek letter $\eta$ (eta) is the learning rate, a small positive number you choose. Updating every knob once is one step.

Why not set $\partial L/\partial w=0$ and solve, the way calculus finds a minimum (Modules 13 and 14)? For one knob you could. For millions of knobs tangled together through $\tanh$, no formula exists. Walking downhill works anyway.

Worked example. $L(w)=(w-3)^2$ has its minimum at $w=3$, and $dL/dw=2(w-3)$. Start at $w_0=0$ with $\eta=0.1$:

To see the pattern, subtract 3 from both sides of the update:

The distance to the minimum is multiplied by $1-2\eta$ at every step. With $\eta=0.1$ that factor is $0.8$: a steady approach. With $\eta=0.5$ it is $0$: the first step lands exactly on 3. With $\eta=0.9$ it is $-0.8$: each step overshoots to the other side, by less each time. With $\eta=1$ it is $-1$: the steps bounce between 0 and 6 forever. With $\eta>1$ the factor’s size exceeds 1, and every step overshoots further than the last.

Gradient descent on $(w-3)^2$ from $w=0$. Too small is slow, too large explodes, and the safe range depends on how sharply the loss curves. With many knobs you find a good $\eta$ by experiment (Module 2 runs a learning-rate sweep).

Two knobs at once. With two knobs the loss is a surface, and a contour map shows it from above: each closed curve joins points of equal loss, like the height lines on a hiking map. The gradient $\nabla L=(\partial L/\partial w_1,\ \partial L/\partial w_2)$ points straight uphill, at right angles to the contour through your position, so the update $w\leftarrow w-\eta\nabla L$ steps straight downhill. (Calculus III proves the right angle; the figure below lets you see it.)

Take $L(w_1,w_2)=\tfrac12\left(w_1^2+8w_2^2\right)$. Its gradient is $(w_1,\ 8w_2)$: the bowl is 8 times steeper in the $w_2$ direction. Each knob gets its own factor, exactly as in the one-knob example: $w_1\leftarrow(1-\eta)\,w_1$ and $w_2\leftarrow(1-8\eta)\,w_2$. The steep direction limits the learning rate: $w_2$ converges only when $|1-8\eta|<1$, that is, $\eta<0.25$. At $\eta=0.23$ the $w_2$ factor is $1-1.84=-0.84$, so the path zigzags across the valley. At $\eta=0.05$ nothing zigzags, but $w_1$ moves only 5% of its remaining distance per step, so progress along the valley floor is slow.

Gradient descent on $L=\tfrac12(w_1^2+8w_2^2)$ from $(-4.5,\,1.2)$. Both axes use the same scale, so the right angle between the downhill arrow and the contour is real. A bowl much steeper in one direction than another is called ill-conditioned. Its steepest direction caps the learning rate, and its gentlest direction then crawls. Optimizers such as Adam (Module 2) give each knob its own effective step size, partly to fix exactly this.

Real losses have many valleys, and gradient descent finds a nearby low point, not necessarily the lowest. For neural networks this works far better than early theory predicted, which is one of the field’s genuine surprises.

Where this goes: every model at Anthropic, OpenAI, and Tesla trains with this loop. The differences are scale: tensors instead of single numbers, GPUs, and smarter step rules such as AdamW (Module 2).

2.7. Checking gradients with finite differences#

How do you know a backward pass is right? Go back to the definition of the derivative, Module 08’s difference quotient, with a small $h$ instead of a limit:

For a partial derivative, nudge one input and freeze the others.

Why central is better, Algebra First. Take $f(x)=x^3$ and expand $(x+h)^3=x^3+3x^2h+3xh^2+h^3$.

• Forward: $\dfrac{(x+h)^3-x^3}{h}=3x^2+3xh+h^2$. The error, $3xh+h^2$, shrinks in proportion to $h$.
• Central: $(x-h)^3=x^3-3x^2h+3xh^2-h^3$, so $(x+h)^3-(x-h)^3=6x^2h+2h^3$. Dividing by $2h$ gives $3x^2+h^2$. The error, $h^2$, shrinks in proportion to $h^2$: the $3xh$ terms cancel because the two nudges are symmetric.

At $x=2$ with $h=0.1$ the true slope is 12. The forward difference gives 12.61 and the central difference 12.01 (AI01.08).

Why not use a tiny $h$? A computer stores about 16 significant digits (a Python float is a 64-bit double, with relative precision about $2.2\times10^{-16}$). When $h$ is tiny, $f(x+h)$ and $f(x-h)$ agree in almost every digit. Subtracting cancels the digits that agree and leaves mostly rounding noise, and dividing by a tiny $2h$ magnifies that noise. So rounding error grows like $1/h$ while the approximation error shrinks like $h$ or $h^2$.

Error in the estimated slope of $\tanh$ at $x=0.5$, both axes logarithmic. On the right, the approximation error dominates (slope 1 for forward, 2 for central). On the left, rounding error dominates and gets worse as $h$ shrinks. The central difference is most accurate near $h=10^{-6}$, with an error around $10^{-12}$; the forward difference is best near $h=10^{-8}$, around $10^{-9}$.

The Gradient Explorer’s third panel draws this curve live in your browser; slide $h$ and watch the error fall, bottom out, and climb again. Practical rule: use the central difference with $h$ between $10^{-5}$ and $10^{-6}$, and compare with a tolerance of about $10^{-5}$ times the size of the gradient. AI01.16 derives the best $h$ with Module 14’s optimization. Gradient checks are to autograd what unit tests are to any other program: too slow for training (two evaluations per input), perfect for tests.

2.8. Python for a C++ engineer: exactly what this lab uses#

2.9. Design trade-offs: why this design, and what breaks at scale#

• Scalars versus tensors. Your engine stores one number per Value, and every operation allocates a Python object and a closure. Measured with a complete engine of this design on one desktop PC, one operation costs about 1.5 to 3 microseconds (forward plus backward; building nodes is the slow, variable part) and about 860 bytes of memory. One $4096\times4096$ matrix multiplication, a routine step in a small language model, is $4096^3\approx6.9\times10^{10}$ multiply-adds, or about $1.4\times10^{11}$ of your nodes. At that rate it is two to five days of computing and 118 terabytes of memory. A GPU does it in a few milliseconds. A tensor library stores millions of numbers per node: one matrix-multiply node replaces millions of your nodes and runs on the GPU. The chain rule and the topological sort are identical. Module 2 makes that move. This scale is for understanding, not speed.
• Memory. The graph keeps every intermediate value alive until the backward pass, because local derivatives need them (a multiply node needs both inputs’ data). That is why training uses far more memory than running a model. Large models trade compute for memory by recomputing some intermediates during the backward pass, which is called activation checkpointing.
• Built as it runs, or compiled first. Your graph is built by ordinary Python as it runs, so any if or for works and an error points at your own line. PyTorch’s default mode works the same way. Compiled approaches such as torch.compile and JAX trace the graph first and then optimize it: faster, harder to debug. Module 2 uses both.
• Reverse mode versus forward mode. One backward pass gives the gradient of one output with respect to every input. Forward mode (W04’s dual numbers) carries the derivative with respect to one input forward per sweep, so it needs one sweep per input. One loss and a billion knobs means reverse mode (AI01.15).
• Accumulating gradients is a feature. Because backward adds into grad, gradients from several backward calls sum. Large-model training uses this to add up gradients over several small batches (“gradient accumulation”). The price is that you must zero the gradients before each step (bug B8). PyTorch does the same for parameters.
• The order of addition changes the last digits. Floating-point addition is not exactly associative: $(a+b)+c$ and $a+(b+c)$ can differ in the last bit, because each sum is rounded. If _prev were a set, Python would order the inputs by memory address, which changes from run to run, so gradients would be summed in a different order each time. Usually that is invisible. In training, though, tiny differences can grow. With a set, XOR runs of a complete engine gave a different final loss on every rerun for all 50 random starts tried; in one batch of five reruns, a single start ranged from $1.3\times10^{-30}$ to $3.7\times10^{-11}$. A tuple fixes the order and makes every run identical, which Module 5’s experiment loop depends on. (GPU libraries face the same issue at a much larger scale, and offer “deterministic” modes that trade speed for identical results.)
• A closure per node, or a table of operations. Storing a closure in each node is short and flexible. PyTorch instead records which operation made each node and calls a compiled backward function for that operation, which is faster and uses less memory per node.

2.10. Beginner bugs: a field guide#

These are the mistakes people really make when they first write an autograd engine. Each has a symptom you can recognize and at least one test in your lab that catches it. Each was planted, one at a time, in a working engine, and every one of these “mutants” fails at least one of the lab’s tests (AI00, section G).

3. The lab: build it yourself#

3.0. Setup and ground rules#

Everything lives in the companion code’s ai-lab/m01_autograd folder; copy the ai-lab folder into a repository of your own first (AI00, section G). You write engine.py, gradcheck.py, nn.py, and train.py; each starts as a set of docstrings that are its specification. The tests folder and check.py come with the starters; don’t edit them. Open a terminal (PowerShell on Windows) in the folder that holds your copy and run the first check:

cd ai-lab\m01_autograd
python check.py 1

Every test fails with “not written yet”. That is the starting line.

• Type the code yourself. Don’t paste from this book or anywhere else: typing is how your hands learn the patterns.
• Predict before every run, in writing. The orange boxes tell you what to predict.
• When a test fails, read the test before your code. Its name says what it checks, and its body in tests/ shows the exact inputs. Reading someone else’s test code is code-reading practice too.
• Stuck? Climb that step’s hint ladder (after the practice set, on its own page) one rung at a time.
• Checking your work. The tests are the judge: a step is finished when its tests pass on code you typed yourself. When one fails, find the first incorrect step before you change anything, and log it with its field-guide number (section 2.10). After all eight steps pass, the code-reading drill (section 4.2) compares your engine with micrograd, a complete published engine of the same design.

3.1. Step 1: a Value that remembers where it came from (worked example)#

Goal. A Value holds a number and a gradient and remembers its inputs; + and * build new Values. Type this into engine.py, replacing the matching stubs:

    def __init__(self, data, _children=(), _op="", label=""):
        self.data = float(data)
        self.grad = 0.0
        self._prev = tuple(_children)
        self._op = _op
        self.label = label
        self._backward = lambda: None

    def __repr__(self):
        return f"Value(data={self.data}, grad={self.grad})"

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        return out

Line by line: float(data) stores a float even when you pass an int. grad starts at 0 because nothing is known yet. _prev holds the inputs, in order, as a tuple. (For a * a it holds a twice; the visited set in step 3 makes sure a is still processed once.) A tuple keeps every run identical to the last digit; section 2.9 explains why a set would not. _op and label are for pictures and debugging. _backward starts as a function that does nothing, which is right for a leaf; every operation will replace it. In __add__, the first line wraps a plain number so that a + 1 works.

Now write __mul__ yourself. It has the same shape as __add__. Only the arithmetic and the operation symbol change.

Run python check.py 1. Passing proves the forward arithmetic, the graph bookkeeping, the wrapping of plain numbers, and that a Value can live in a set.

3.2. Step 2: local backward functions (partly filled)#

Goal. Each operation attaches a _backward that multiplies its local derivatives by out.grad and adds the results into its inputs’ gradients. Extend __add__ and fill in the two blanks from the table in section 2.3:

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")

        def _backward():
            self.grad += ____ * out.grad    # d(out)/d(self) for out = self + other
            other.grad += ____ * out.grad   # d(out)/d(other)
        out._backward = _backward
        return out

_backward is defined inside __add__ so that it is a closure: it remembers self, other, and out, and it runs later, once out.grad is known. Now give __mul__ its _backward on your own. Then run the backward pass by hand in the Python prompt:

from engine import Value
a, b, c = Value(3.0, label="a"), Value(-2.0, label="b"), Value(5.0, label="c")
e = a * b; d = e + b; L = d * c
L.grad = 1.0
L._backward(); d._backward(); e._backward()
print(a.grad, b.grad, c.grad)

Run python check.py 2.

3.3. Step 3: backward() for the whole graph (on your own, from the recipe)#

Goal. L.backward() fills in every gradient in one call. No code this time: use the recipe in section 2.5 and the three numbered steps in the docstring (sort, seed, walk in reverse). Use a set of visited nodes, and let the output’s gradient with respect to itself be 1.

Run python check.py 3. Then watch it work with the companion code’s visualizer, from the same folder:

python ..\tools\viz.py example

This builds the running example with your engine and opens out/graph.html, an animation of the forward pass and then of backprop, one node at a time. It replays your own _backward functions. Step through and watch $b$ receive $+5$ when $d$ reports and $+15$ when $e$ reports. If a gradient looks wrong in the animation, the frame where it goes wrong names the operation to fix.

3.4. Step 4: break it on purpose#

This step is an experiment, not new code.

1. In __add__’s _backward, change both += to =. Predict a.grad after (a + a).backward() with $a=3$. Run python check.py 4 and note which tests fail. Restore the +=.
2. Do the same in __mul__ and predict a.grad for (a * a).backward() at $a=3$. Run, note, restore.
3. Write one sentence in your log: why does = give exactly 1 for $a+a$?

Run python check.py 4 once more with everything restored; all four steps should pass.

3.5. Step 5: the remaining operations (on your own)#

Goal. __pow__, __neg__, __radd__, __sub__, __rsub__, __rmul__, __truediv__, __rtruediv__, exp, tanh, and relu. Their docstrings give the contracts. Before typing each one, write its local derivative on paper (the Workbench).

Build from what you have. An operation made out of existing operations needs no _backward of its own, because the pieces already have one. For example $-x=x\cdot(-1)$, $a-b=a+(-b)$, and $a/b=a\cdot b^{-1}$. The trade-off: more nodes (slower) in exchange for less code and fewer places for bugs. __pow__, exp, tanh, and relu need their own _backward. If your power rule drops the sign of the new exponent, check.py reports -48.0 != -0.75 in test_pow_negative_exponent. That $-48$ is not a coincidence you should shrug at: it is the test’s $3\cdot(-2)\cdot2^{3}$ where $3\cdot(-2)\cdot2^{-3}$ belongs, the same slip as answering warm-up question 1 with $-48x^4$.

Run python check.py 5.

3.6. Step 6: a gradient checker (on your own)#

Goal. numerical_grad and max_grad_error in gradcheck.py, following section 2.7 and the docstrings. Step 6 also includes a test that builds 40 random expressions from all of your operations and compares your backward pass with central differences computed by the test itself. It is the strongest single check in the lab.

Run python check.py 6.

3.7. Step 7: neurons, layers, networks (partly filled)#

MLP(2, [4, 4, 1]). Every line is a weight. Each neuron also has a bias. Hidden layers apply $\tanh$; the last layer does not, so the output can be any number.

Goal. nn.py: a Neuron computes $\tanh(w\cdot x+b)$ (or the plain sum when nonlin=False); a Layer is a list of neurons reading the same inputs; an MLP (multilayer perceptron) applies layers one after another with a linear last layer. Everything is ordinary Python on top of Value, with no new calculus. Here is the neuron with its two key lines left blank:

class Neuron(Module):
    def __init__(self, nin, nonlin=True):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(nin)]
        self.b = Value(0.0)
        self.nonlin = nonlin

    def __call__(self, x):
        act = ____        # w0*x0 + w1*x1 + ... + b, as one Value
        return ____       # tanh of it, unless nonlin is False

Then write Neuron.parameters, Module, Layer, and MLP yourself.

Why random starting weights? If every weight started equal, every neuron in a layer would compute the same thing and receive the same gradient, forever. Random starts break that symmetry. The bias can start at 0 because the weights already make the neurons different. Why a linear last layer? So the output is not trapped between $-1$ and $1$; the loss decides what scale it needs.

Run python check.py 7.

3.8. Step 8: teach it XOR (on your own)#

XOR’s output is $+1$ (blue) when exactly one input is 1, and $-1$ (red) otherwise.

Why XOR needs a hidden layer. Left: one neuron is blue on one side of a straight line and red on the other, and no straight line can put both blue points on one side and both red points on the other. Right: a network wired by hand with two hidden neurons, $h_1=\tanh\bigl(4(x_1+x_2-0.5)\bigr)$ and $h_2=\tanh\bigl(4(x_1+x_2-1.5)\bigr)$, and output $h_1-h_2-1$. It gives $-0.96$, $+0.93$, $+0.93$, and $-0.96$ at the four corners.

Why a hidden layer works. A $\tanh$ neuron outputs nearly $+1$ on one side of the line $w\cdot x+b=0$ and nearly $-1$ on the other, so each neuron draws one soft straight line. On its own it can only split the square into two half-planes. Two hidden neurons draw two lines, and the output neuron adds their votes, so it can keep just the band between them. Training finds weights like the hand-wired ones by itself. Stacking layers this way, building complicated regions out of simple pieces, is the core idea behind deep networks.

Goal. In train.py: mse_loss (the mean of the squared errors), train_step (zero the gradients, forward, backward, update, in that order, returning the loss as a float), and train (repeat, collecting the losses).

Run python check.py 8. Then train and look:

python train.py
python train.py --viz

The second command opens out/xor.html: your loss curve and a map of what your network outputs over the whole input square.

A complete solution’s run with seed 1: MLP(2, [4, 4, 1]), learning rate 0.1, 400 steps of full-batch gradient descent (every step uses all four examples). Left: the loss on a logarithmic scale. It keeps falling at a steady rate (a straight line on this scale) because a network with a linear output can fit four points exactly; on real data the loss levels off instead. Right: the trained network’s output over the input square (blue near $+1$, red near $-1$, pale near 0), with the four training points. Your curve will differ in detail, because random starting weights depend on how your code draws them.

The spike is not a bug. In that run the loss falls to about $7.7\times10^{-5}$ by step 74, then shoots up to 2.87 at step 84, higher than where it started, and then falls again. Near the bottom of this valley the loss curves much more sharply than it did at the start. A learning rate that was safe on the gentle early slope is now too big, which is section 2.6’s $\eta>1$ bounce, happening in 37 dimensions at once. The overshoot throws the weights somewhere flatter, and descent resumes. In runs of a complete solution at learning rate 0.1, 16 of 50 random starts climb back above a loss of 1 after first falling below 0.01 (the worst reaches 6.9), and all 50 still learn XOR. If your curve spikes, check the loss on the steps after the spike before suspecting your code; a smaller learning rate such as 0.05 removes the spikes but learns more slowly.

This is not a toy curiosity; researchers study it directly. As training goes on, the loss surface around the weights often gets sharper (progressive sharpening), and once its sharpest curvature passes $2/\eta$ ($=20$ here), a gradient-descent step overshoots, exactly as in the one-knob example. Cohen et al. (2021) documented this “edge of stability” in full-batch gradient descent. A blow-up that ends in a flatter region also resembles what Lewkowycz et al. (2020) call a catapult, which they observed at the very start of training with a large learning rate. Your spike is consistent with this picture, and stretch S6 shows how you could test it.

When python check.py reports all eight steps passing, you have written a working autograd engine and trained a neural network with it. Copy the summary line into the development log (section 5).

4. Practice, code reading, and the gate#

4.1. Practice set#

F = foundation; C = standard; T = transfer; S = optional stretch; L = lab; R = code reading; G = mastery gate. Follow the session plan in section 1. Attempt each problem completely before opening its solution, and log every miss in error_log.csv with its friction family.

4.2. Code-reading drill: micrograd#

micrograd, by Andrej Karpathy (MIT license), is the classic hundred-line autograd engine. The engine you built follows the same design, which makes it the gentlest possible first read of someone else’s code. Read the file at this pinned commit; the line numbers below refer to it:

Protocol. Read the whole file once, top to bottom, without stopping (five minutes). Then, for each question, write your explanation in your own words before opening its solution. Describe what the code does, not what you think it should do.

4.3. Mastery gate#

# inside __mul__
def _backward():
    self.grad = other.data * out.grad
    other.grad = self.data * out.grad

# inside tanh, where t = math.tanh(self.data)
def _backward():
    self.grad += (1 - t) * out.grad