Autograd

PyTorch’s autograd is a reverse-mode automatic differentiation engine that records operations into a graph as they run, then walks that graph backward to compute gradients. This is the concrete machinery behind Backpropagation.

The dynamic graph is built during the forward pass

There is no separate “compile the graph” step in eager mode. Every time you run a tensor op on an input that needs gradients, autograd attaches a node describing how to reverse that op. Run different code (an if, a Python loop of variable length) and you get a different graph. This is why PyTorch is called define-by-run.

Two attributes drive everything:

• Tensor.requires_grad: a leaf flag. Set it on the inputs (parameters) you want gradients for.
• Tensor.grad_fn: on a non-leaf output, this points to the Node that knows how to compute the backward of the op that produced it.

import torch
 
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = (x * x).sum()          # y.grad_fn -> SumBackward0
print(y.grad_fn)           # <SumBackward0 object at ...>
print(y.grad_fn.next_functions)  # edges back to MulBackward0, then to x

Leaves vs intermediates

A leaf is a tensor you created directly with requires_grad=True (typically parameters). Only leaves accumulate into .grad by default. Intermediates have a grad_fn but their .grad is None unless you call .retain_grad().

Function / Node: forward and backward

Each differentiable op is a torch.autograd.Function with two static methods. forward computes the output and stashes whatever the backward needs in ctx. backward receives the gradient of the loss w.r.t. this op’s output (grad_output) and returns the gradient w.r.t. each input, applying the chain rule.

forward:  inputs --op--> output        (also: ctx.save_for_backward(...))
backward: grad_output --VJP--> grad_inputs

Each backward computes a vector-Jacobian product (VJP): it never materializes the full Jacobian, it just left-multiplies grad_output by it.

The backward pass and accumulation

loss.backward() seeds the output gradient with 1.0 (for a scalar) and does a reverse topological traversal, calling each node’s backward and routing results along the next_functions edges. When two paths reach the same leaf, the gradients add. This accumulation is why you must zero grads between steps:

optimizer.zero_grad(set_to_none=True)   # else this step's grads add to last step's
loss.backward()
optimizer.step()

Accumulation is a feature, not a bug

Gradient accumulation across micro-batches relies on this: call backward() several times before step() to simulate a larger batch. See PyTorch and Memory for how activations dominate the memory cost of the graph.

Turning autograd off: no_grad and inference_mode

Recording the graph costs memory (saved tensors) and time. During evaluation you do not need it.

with torch.no_grad():           # do not record; outputs have requires_grad=False
    logits = model(x)
 
with torch.inference_mode():    # stronger: also skips version counters / view tracking
    logits = model(x)

inference_mode is faster but its outputs cannot later be used in autograd, so use no_grad if you might need to re-enable grad on the results.

Custom autograd.Function

Write your own when you have a fused kernel, a numerically better backward, or a non-differentiable step you want to give a custom gradient. Here is a numerically stable LogSumExp:

class LogSumExp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, dim):
        m = x.max(dim=dim, keepdim=True).values
        out = (x - m).exp().sum(dim=dim, keepdim=True).log() + m
        ctx.save_for_backward(x, out)
        ctx.dim = dim
        return out.squeeze(dim)
 
    @staticmethod
    def backward(ctx, grad_output):
        x, out = ctx.saved_tensors
        # softmax along dim is the local gradient; broadcast grad back
        softmax = (x - out).exp()
        return grad_output.unsqueeze(ctx.dim) * softmax, None  # None for `dim`
 
y = LogSumExp.apply(torch.randn(4, 8, requires_grad=True), 1)

Check yourself with gradcheck

torch.autograd.gradcheck(fn, inputs) compares your analytic backward against finite differences (use float64 inputs). It catches sign errors and missing terms immediately.

A tiny reverse-mode engine from scratch

To demystify grad_fn, here is a minimal scalar autograd. Each Value stores its data, a grad, and a _backward closure capturing its local derivative.

class Value:
    def __init__(self, data, _children=()):
        self.data = data
        self.grad = 0.0
        self._backward = lambda: None
        self._prev = set(_children)
 
    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other))
        def _backward():                 # d(out)/d(self) = other.data, etc.
            self.grad  += other.data * out.grad
            other.grad += self.data  * out.grad
        out._backward = _backward
        return out
 
    def __add__(self, other):
        out = Value(self.data + other.data, (self, other))
        def _backward():
            self.grad  += out.grad       # gradient flows through unchanged
            other.grad += out.grad
        out._backward = _backward
        return out
 
    def backward(self):
        topo, seen = [], set()
        def build(v):                    # reverse topological order
            if v not in seen:
                seen.add(v)
                for c in v._prev: build(c)
                topo.append(v)
        build(self)
        self.grad = 1.0                  # seed dL/dL = 1
        for v in reversed(topo):
            v._backward()
 
a, b = Value(2.0), Value(-3.0)
L = a * b + a                            # L = a*b + a
L.backward()
print(a.grad, b.grad)                    # 2.0, ... ; dL/da = b + 1 = -2.0

This is exactly what PyTorch does, scaled up: _backward closures are the Node.backward methods, _prev edges are next_functions, and topo is the traversal engine.cpp performs.

Higher-order gradients

Because backward is itself built from differentiable ops, you can differentiate the gradient. Pass create_graph=True so the backward pass records its own graph.

x = torch.tensor(1.5, requires_grad=True)
y = x ** 3
g, = torch.autograd.grad(y, x, create_graph=True)   # dy/dx = 3x^2
g2, = torch.autograd.grad(g, x)                      # d2y/dx2 = 6x
print(g.item(), g2.item())                           # 6.75, 9.0

This powers Hessian-vector products, meta-learning (MAML), and gradient penalties (WGAN-GP). It is memory-hungry: the first backward graph is retained. See Tensors, Storage, and Strides for how saved tensors are stored, and torch.compile and Inductor for how AOTAutograd captures the backward as a traceable graph.