The Computational Graph and grad_fn

Every operation builds a node

When you operate on a tensor with requires_grad=True, PyTorch creates a graph node that knows two things: what operation produced this tensor (its grad_fn) and which inputs fed into the operation (so backward can walk back to them).

You can inspect this directly:

Leaf tensors — created directly (e.g. nn.Parameter, or any torch.tensor(..., requires_grad=True)). Their grad_fn is None and is_leaf is True. They are where gradients land.
Intermediate tensors — produced by operations. Their grad_fn points to a backward function (AddBackward0, MulBackward0, etc.) and they are not leaves. By default they don't store their own .grad — only leaves do.

Walking the graph backwards

When you call loss.backward(), autograd traces from the loss back through all the connected nodes via the grad_fn links. Each node knows how to compute the local Jacobian (or vector-Jacobian product, technically), and the chain rule composes them into the final gradient at every leaf.

One subtlety: by default the graph is freed immediately after backward to save memory. If you need to call backward again on the same graph, pass retain_graph=True. You usually don't.

Code

Inspecting the graph·python

import torch

x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

z = w * x          # MulBackward0
y = z + b          # AddBackward0
loss = y ** 2      # PowBackward0

# Leaves
print(x.is_leaf, x.grad_fn)   # True None
print(w.is_leaf, w.grad_fn)   # True None

# Intermediates
print(z.is_leaf, z.grad_fn)   # False <MulBackward0>
print(y.is_leaf, y.grad_fn)   # False <AddBackward0>
print(loss.grad_fn)           # <PowBackward0>

Backward through a tiny linear-regression-like graph·python

import torch

x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

z = w * x          # 6
y = z + b          # 7
loss = y ** 2      # 49
loss.backward()

# Gradients via chain rule
# dloss/dy = 2y       = 14
# dy/dz   = 1         → dloss/dz = 14
# dy/db   = 1         → dloss/db = 14
# dz/dw   = x = 2     → dloss/dw = 28
# dz/dx   = w = 3     → dloss/dx = 42
print(x.grad, w.grad, b.grad)
# tensor(42.) tensor(28.) tensor(14.)

retain_graph — when you need backward twice·python

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3

# First backward — frees graph by default
y.backward()
print(x.grad)         # 12.

# Second backward without retain_graph errors:
# y = x ** 3
# y.backward()  -- works because we built a fresh graph

# If you wanted to backward TWICE on the SAME graph:
x.grad = None
y = x ** 3
y.backward(retain_graph=True)
y.backward()          # second pass on the same retained graph
print(x.grad)         # 24. (gradients accumulate)

The Computational Graph and grad_fn

Every operation builds a node

Walking the graph backwards

Code

External links

Exercise

Progress

Comments 0