Autograd and computational graphs in PyTorch: Visual walkthrough

This page rebuilds the parent result — autograd computing $\frac{\partial z}{\partial x}$ and $\frac{\partial z}{\partial y}$ for $z = x y + x^{2}$ — but as a sequence of pictures. Every symbol is earned before it is used. If you have never seen a "gradient" or a "graph node", start here at line one.

We use only two starting numbers, the ones from the parent note: $x = 2, y = 3.$

Step 1 — What is a computational graph, really?

WHAT. Every arithmetic expression is secretly a little assembly line. Write $z = x y + x^{2}$ as three simple machines wired together:

a multiply machine that eats $x$ and $y$ and outputs $a = x y$ ,
a square machine that eats $x$ and outputs $b = x^{2}$ ,
an add machine that eats $a$ and $b$ and outputs $z = a + b$ .

WHY. The reason we bother splitting one formula into tiny machines is that each tiny machine has a derivative we already know by heart (multiply, square, add). If we can chain those known derivatives, we never have to differentiate the whole messy formula at once. That chaining is what the graph is for.

PICTURE. Circles are tensors (the numbers flowing), rounded boxes are operations (the machines). Arrows show which number feeds which machine.

Figure — Autograd and computational graphs in PyTorch

Notice $x$ has two arrows leaving it: one into the multiply machine, one into the square machine. Hold onto that — it is the whole reason gradients get summed later.

Step 2 — The forward pass: push the numbers through

WHAT. Plug in $x = 2, y = 3$ and let each machine fire in order.

$a = x \cdot y = 2 \cdot 3 = 6, b = x^{2} = 2^{2} = 4, z = a + b = 6 + 4 = 10.$

Term by term: $a$ is the multiply-box output, $b$ is the square-box output, $z$ is their sum — the final answer sitting at the end of the line.

WHY. We must run forward first because every machine needs to remember the actual numbers that passed through it. The square machine, for instance, will later need to recall "I saw $x = 2$ " to report how sensitive it is. No forward values → no backward gradients.

PICTURE. Same graph, now every arrow is labelled with the real number riding along it.

Step 3 — What does a single machine's derivative mean?

WHAT. Before wiring the whole backward pass, look at ONE machine. Take the square machine, $b = x^{2}$ . Its derivative is $\frac{\partial b}{\partial x} = 2 x .$

The symbol $\frac{\partial b}{\partial x}$ (read "partial dee $b$ dee $x$ ") answers exactly one question: if I nudge the input $x$ up by a tiny amount, how many times bigger is the wobble that comes out at $b$ ?

WHY this tool — the derivative and not something else. We want sensitivity. "How much does the output move per unit of input move?" is precisely what a derivative measures. No other tool answers "rate of change at a point". At $x = 2$ the answer is $2 \cdot 2 = 4$ : a tiny push on $x$ produces a wobble $4 \times$ as big on $b$ .

PICTURE. The curve $b = x^{2}$ with a tiny input step $h$ at $x = 2$ and the much bigger output step it causes. The slope of the tangent line is the number 4.

Step 4 — The chain rule: multiply the slopes along a path

WHAT. Ask: how sensitive is the final $z$ to the square-machine's input, following the path $x \to b \to z$ ? The chain rule says: multiply the local slopes along the path. $along path x \to b \to z \frac{\partial z}{\partial x} = how z reacts to b \frac{\partial z}{\partial b} \cdot how b reacts to x \frac{\partial b}{\partial x} = 1 \cdot 2 x .$

WHY multiplication? If wiggling $x$ makes $b$ move $2 x$ times as much, and wiggling $b$ makes $z$ move $1$ time as much, then wiggling $x$ makes $z$ move $2 x \cdot 1$ times as much. Sensitivities compound — that is multiplication, exactly like gears: two gears with ratios $r_{1}$ and $r_{2}$ give a total ratio $r_{1} r_{2}$ .

PICTURE. A gear train: input $x$ turns a gear, which turns $b$ , which turns $z$ . The total turn ratio is the product of the two gear ratios written on the edges.

Step 5 — The node $x$ has two paths: SUM them

WHAT. But $x$ also reaches $z$ by a second road, $x \to a \to z$ . Its sensitivity along that road is $\frac{\partial z}{\partial a} \cdot \frac{\partial a}{\partial x} = 1 \cdot y .$ The multivariate chain rule says the total sensitivity of $z$ to $x$ is the sum over every path: $\frac{\partial z}{\partial x} = road through a 1 \cdot y + road through b 1 \cdot 2 x = y + 2 x .$

WHY sum? Wiggling $x$ shoves $z$ through both roads at once, and the two shoves land at $z$ together. Independent contributions to the same quantity add. (This is the "sum over children" rule quoted in the parent's chain-rule callout.)

PICTURE. Two coloured arrows leaving $x$ , each carrying its own contribution ( $y$ in green, $2 x$ in yellow), meeting and adding at $x$ 's gradient slot.

Plugging in $x = 2, y = 3$ : $\frac{\partial z}{\partial x} = 3 + 2 (2) = 3 + 4 = 7.$

And $y$ has only one road, $y \to a \to z$ : $\frac{\partial z}{\partial y} = \frac{\partial z}{\partial a} \cdot \frac{\partial a}{\partial y} = 1 \cdot x = 2.$

Step 6 — The backward pass as water flowing back

WHAT. PyTorch does not think in "paths". It floods the graph backward. It plants a $1$ at $z$ (because $\frac{\partial z}{\partial z} = 1$ : $z$ is perfectly sensitive to itself), then at every machine it multiplies the incoming backward number by that machine's local slope and pushes the result upstream. Where two streams meet a node, they add.

$start: \frac{\partial z}{\partial z} = 1 \to \frac{\partial z}{\partial a} = 1, \frac{\partial z}{\partial b} = 1 \to merge at x : 1 \cdot y + 1 \cdot 2 x = 7.$

WHY do it backward instead of forward? Because we want the gradient with respect to many inputs but of one scalar output. Flooding backward from the single output visits every input once. Flooding forward would need a separate sweep per input — far more work. This is the core efficiency of reverse-mode differentiation, the backpropagation algorithm.

PICTURE. The graph with backward arrows in red; each edge labelled by the number travelling back, ending at $x . grad = 7$ and $y . grad = 2$ .

Step 7 — The degenerate case: what if a machine's slope is zero?

WHAT. Suppose one input never affects the output, like $y$ inside the square machine $b = x^{2}$ : $\frac{\partial b}{\partial y} = 0$ . Or a ReLU that saw a negative number: its local slope is $0$ .

WHY it matters. A local slope of $0$ acts like a closed valve: multiply any backward stream by $0$ and nothing passes upstream. That input receives no gradient through that path. This is not a bug — it is the machine honestly reporting "I ignore this input, so nudging it changes nothing".

PICTURE. Two backward valves side by side: an open valve (slope $1$ , water flows) and a shut valve (slope $0$ , water stops), showing a dead path where gradient vanishes.

Step 8 — The other edge case: two backward visits accumulate

WHAT. If you press .backward() twice without clearing, the second flood adds onto the first, because .grad is a running total, not a fresh slate: $x . grad \leftarrow x . grad + \frac{\partial z}{\partial x} .$ So after two backward passes on $y = x^{2}$ at $x = 2$ : first gives $4$ , second gives $4 + 4 = 8$ .

WHY. Mini-batch training wants a sum: the batch gradient is $\frac{1}{N} \sum_{i} \frac{\partial L _{i}}{\partial w}$ , so accumulating each sample's gradient builds the sum for free. The price: you must call .zero_() between real steps, or old gradients leak into new ones.

PICTURE. A bucket labelled .grad filling up: first pour adds $4$ , second pour adds another $4$ → level $8$ . A tap labelled zero_() empties it.

The one-picture summary

Everything above, compressed: forward numbers in blue, local slopes on the boxes, red backward flow merging at $x$ to give $7$ and reaching $y$ to give $2$ .

Recall Feynman retelling — say it back in plain words

I wrote my formula as a little assembly line of tiny machines: multiply, square, add. I pushed my numbers $2$ and $3$ forward and each machine remembered what it saw. Then I planted a single "1" at the very end — because the answer is perfectly sensitive to itself — and let it flow backward. At each machine I multiplied the incoming backward number by that machine's own slope (multiply's slope is the other input, square's slope is $2 x$ , add's slope is just $1$ ). Where $x$ 's two roads met, I added their contributions, getting $y + 2 x = 3 + 4 = 7$ ; $y$ had one road and got $x = 2$ . If a machine's slope were $0$ , that road would be a shut valve and no gradient would pass. And if I run it backward twice without emptying the bucket, the gradients pile up — which is a feature for batches but a trap if I forget to zero them.

Recall

In $z = x y + x^{2}$ , why does $x$ get gradient $y + 2 x$ but $y$ only gets $x$ ? ::: $x$ reaches $z$ by two paths (through $x y$ and through $x^{2}$ ), and their contributions sum; $y$ reaches $z$ by only one path (through $x y$ ). What single number do we plant at the output to start the backward pass, and why? ::: We plant $1$ , because $\frac{\partial z}{\partial z} = 1$ — the output is perfectly sensitive to itself. What does a local slope of $0$ (e.g. a negative-input ReLU) do to backward flow? ::: It shuts the valve — any gradient multiplied by $0$ vanishes, so no gradient passes upstream through that path. Why must you call .zero_() between training steps? ::: .grad accumulates (adds) each backward pass; without zeroing, old gradients leak into the next step's update.

See also: tensors · backpropagation · torch.nn.Module · gradient checkpointing (trades memory for recomputing this graph).