Frameworks and Automatic Differentiation

A machine-learning framework defines the executable contract between model code and hardware. That contract covers tensor values, operator semantics, derivative rules, device placement, distributed layout, and execution mode. A program can be mathematically sound and still fail because one of these parts is undefined, unsupported, or interpreted differently by another backend.

This chapter follows that contract from a scalar loss to the work performed on devices. It first explains automatic differentiation, then the memory and correctness obligations of reverse mode, and finally the framework machinery that captures, compiles, dispatches, and distributes the resulting program. The hardware limits underneath it are introduced in Chapter 62; the compiler and kernel pipeline continues in Chapter 64.

Three ways to obtain a derivative

Finite differences estimate a derivative by perturbing an input. For a scalar function $f$ evaluated at $x$, the forward-difference estimate for coordinate $i$ is

$$g_i^{(h)}=\frac{f(x+h e_i)-f(x)}{h} =\frac{\partial f(x)}{\partial x_i} +O(h)+O\!\left(\frac{\epsilon_{\mathrm{mach}}}{h}\right).$$

Here, $x\in\mathbb{R}^n$ is the input, $e_i$ is the basis vector whose $i$th entry is one, $h$ is the nonzero step size, and $g_i^{(h)}$ is the estimate. The $O(h)$ term represents truncation error, while the $O(\epsilon_{\mathrm{mach}}/h)$ term represents scale-dependent floating-point round-off and cancellation. A smaller step reduces one error and can amplify the other. Computing every coordinate also requires one perturbation per input for forward differences, or two for central differences. That is unsuitable for producing a billion-parameter training gradient, although directional finite differences remain valuable as a gradient check.

Symbolic differentiation transforms a mathematical expression into another expression. It can produce a useful closed form, but naive expansion can repeat shared subexpressions and cause expression swell. Control flow, mutation, calls into external libraries, and large tensor programs also require a richer program representation than an ordinary computer-algebra expression.

automatic differentiation takes a different route. It composes a local derivative rule for each primitive in an executed numerical program while preserving the program's intermediate structure (Baydin et al. 2018). The implementation may record an eager tape, trace a graph, transform an intermediate representation, or rewrite source code. In ideal arithmetic it applies the chain rule exactly to the selected differentiable path; on a machine it evaluates that derivative at machine precision and still incurs rounding, overflow, and underflow. It does not differentiate a discrete branch decision. At a nondifferentiable point, the framework must choose a convention, return an undefined value, or raise an error.

The history predates neural-network libraries. Wengert described a linear list of elementary evaluations in 1964 (Wengert 1964). Linnainmaa developed reverse accumulation for computer programs, and the method was later popularized for neural networks as backpropagation (Linnainmaa 1976; Rumelhart et al. 1986). Modern frameworks generalize that idea from scalar arithmetic to tensor operators.

Linearization is the core interface

Consider a differentiable function

$$f:\mathbb{R}^n\to\mathbb{R}^m.$$

Here, $n$ is the input dimension, $m$ is the output dimension, and $x\in\mathbb{R}^n$ is the point at which the function is evaluated. The Jacobian $J_f(x)\in\mathbb{R}^{m\times n}$ contains every local sensitivity, with entry $(j,i)$ equal to $\partial f_j(x)/\partial x_i$. Frameworks rarely need to materialize that whole matrix. They need its action on a vector.

Forward mode computes a Jacobian-vector product (JVP),

$$\operatorname{JVP}_f(x;v)=J_f(x)v\in\mathbb{R}^m,$$

where $v\in\mathbb{R}^n$ is an input tangent. A basis tangent $v=e_i$ returns one column of the Jacobian, while an arbitrary $v$ returns a directional derivative. Forward mode propagates the tangent alongside each intermediate and does not need to retain a reverse tape.

Reverse mode computes a vector-Jacobian product (VJP),

$$\operatorname{VJP}_f(x;u)=J_f(x)^\mathsf{T}u\in\mathbb{R}^n,$$

where $u\in\mathbb{R}^m$ is an output cotangent and $\mathsf{T}$ denotes transpose. A basis cotangent returns one Jacobian row, transposed. If $L:\mathbb{R}^n\to\mathbb{R}$ is a scalar loss, then $m=1$ and the output seed $u=1$ produces the entire gradient in one reverse sweep:

$$\nabla L(x)=J_L(x)^\mathsf{T}1.$$

This geometry determines the mode choice. A full Jacobian can be assembled from $n$ JVP columns or $m$ VJP rows, unless batching or structure reduces the work. Neural-network training has many inputs and one scalar loss, so reverse mode is the natural fit. Forward mode remains useful for directional derivatives, functions with few inputs and many outputs, and compositions such as Hessian-vector products (JAX Authors 2026).

The cheap-gradient result belongs to an arithmetic model, not a wall-clock service-level objective. Baur and Strassen proved a constant-factor bound for a rational straight-line computation and all of its first partial derivatives under a particular operation-count convention (Baur and Strassen 1983). More generally, when primitive VJP rules have bounded relative cost, reverse arithmetic is $O(C_f)$ for a scalar function whose evaluation costs $C_f$. The bound does not include saved-tensor traffic, kernel launches, communication, synchronization, or the storage and writing of $n$ gradient values. Real backward time need not be a fixed multiple of forward time.

How reverse accumulation works

An eager engine records a directed acyclic graph (DAG) of the tensor operations that ran. Each node holds an output value and a backward rule that maps an incoming output adjoint to contributions for its parents. In scalar notation, if node $v_i$ is consumed by later nodes, its accumulated adjoint is

$$\bar v_i=\sum_{j\in\operatorname{succ}(i)} \bar v_j\frac{\partial v_j}{\partial v_i}.$$

Here, $\bar v_i=\partial L/\partial v_i$ is the loss sensitivity of node $v_i$, $\operatorname{succ}(i)$ denotes its direct successors, and the sum collects one contribution from every use. Each node accumulates those contributions. The scalar output's output adjoint is seeded with one. Nodes are then processed in reverse topological order, so all downstream contributions have arrived before a node sends gradients to its parents. Vector-valued operators apply the transpose of a local Jacobian rather than a scalar derivative.

forward:
  execute each primitive and record its parents plus its local VJP rule

reverse:
  set every adjoint to zero and set the scalar output adjoint to one
  visit nodes once in reverse topological order
  add each node's VJP contribution to every parent adjoint

For $v_1=xy$, $v_2=\sin(x)$, and $z=v_1+v_2$, the backward sweep gives $\bar z=1$, $\bar v_1=\bar v_2=1$, $\bar x=\bar v_1y+\bar v_2\cos(x)=y+\cos(x)$, and $\bar y=\bar v_1x=x$. The two contributions to $x$ must be added.

The following runnable implements the scheduling invariant directly. The first case deliberately reuses a non-leaf node: $a=x^2$ and $z=a+a$. A recursive backward routine that propagates an already accumulated adjoint once per path returns the wrong answer for this graph. Topological scheduling processes each node once and returns $\mathrm{d}z/\mathrm{d}x=4x=8$ at $x=2$.

import math

class Var:
    def __init__(self, val, parents=()):
        self.val = float(val)
        self.parents = tuple(parents)  # (parent, local derivative)
        self.grad = 0.0

    def __add__(self, other):
        return Var(self.val + other.val, ((self, 1.0), (other, 1.0)))

    def __mul__(self, other):
        return Var(self.val * other.val, ((self, other.val), (other, self.val)))

    def backward(self, seed=1.0):
        order = topo(self)
        for node in order:
            node.grad = 0.0
        self.grad = seed
        for node in reversed(order):
            for parent, local in node.parents:
                parent.grad += node.grad * local

def sin(value):
    return Var(math.sin(value.val), ((value, math.cos(value.val)),))

def topo(root):
    order, seen = [], set()
    def visit(node):
        if node in seen:
            return
        seen.add(node)
        for parent, _ in node.parents:
            visit(parent)
        order.append(node)
    visit(root)
    return order

def finite_difference(fn, x, h=1e-6):
    return (fn(x + h) - fn(x - h)) / (2 * h)

x = Var(2.0)
a = x * x
z = a + a
z.backward()
print("shared analytic", x.grad, "expected", 8.0)

x = Var(2.0)
z = x * x + sin(x)
z.backward()
numeric = finite_difference(lambda t: t * t + math.sin(t), 2.0)
print("finite difference", round(numeric, 6), "analytic", round(x.grad, 6))

Real engines store vector VJP functions and only the saved tensors needed by those functions; they do not store a scalar local derivative on every edge. They also distinguish two kinds of accumulation. Fan-out contributions are summed inside one backward graph. Separately, parameter .grad buffers often persist across backward calls, so an optimizer loop clears them between steps (PyTorch Contributors 2026).

Differentiability is part of the operator contract

Automatic differentiation is only as correct as the primitive rules it composes. Several cases require an explicit decision:

• A nondifferentiable operation may use a documented subgradient, return zero, propagate a NaN, or reject the request. Ties in max, clipping boundaries, indices, sorting, and integer outputs deserve tests rather than assumptions.
• A detach or stop-gradient operation intentionally cuts a path. Accidentally copying a tensor through an API that drops history can do the same thing.
• In-place mutation or an incorrect alias declaration can overwrite a value saved for backward. Framework version counters catch many cases, but a custom operator still has to declare mutation and aliasing honestly.
• A custom derivative can connect external code to autodiff, but its backward, JVP, batching, autocast, tracing, and higher-order behavior are separate contracts. If an operation can be expressed with built-in tensor operators, composition usually preserves more of those contracts automatically (PyTorch Contributors 2026).
• Mixed precision changes both value and gradient numerics. FP16 training may need loss scaling; gradients must be unscaled before clipping or inspection, and NaN or infinity detection belongs in the optimizer protocol.

A gradient check compares an analytical directional derivative with a centered finite difference in double precision, away from discontinuities. For a high-dimensional input, choose a direction $r$ and compare $\nabla L(x)^\mathsf{T}r$ with $[L(x+hr)-L(x-hr)]/(2h)$. Here, $r$ is the test direction and $h$ is the finite difference step. This checks the whole gradient along one direction without constructing a coordinate-wise numerical gradient. Custom rules should also be checked at zero-size, broadcast, noncontiguous, and extreme inputs, and should receive a second-order gradient check if higher derivatives are promised (PyTorch Contributors 2026).

Reverse mode trades memory for time

The forward pass saves tensors required by later VJP rules. It does not necessarily retain every intermediate, and saved tensors can be released as backward consumes them. Even so, these activations can dominate dynamic graph metadata. Total training memory also includes parameters, gradients, optimizer state, temporary workspaces, communication buffers, and allocator reserves.

Activation checkpointing omits selected saved tensors and reruns part of the forward computation during backward. For a uniform chain of $n$ stages divided into segments of $k$ stages, a simple model for the number of simultaneously stored activation states is

$$M(k)\approx\left\lceil\frac{n}{k}\right\rceil+k.$$

Here, $M(k)$ is the peak saved-state count, $n$ is the chain length, $k$ is the segment length, and $\lceil\cdot\rceil$ rounds upward. Choosing $k\approx\sqrt{n}$ gives $M(k)=O(\!\sqrt{n})$ while recomputing each segment at most once (Chen et al. 2016). Recursive schedules can reach different time-memory points. These bounds do not transfer unchanged to arbitrary DAGs, unequal activation sizes, skip connections, or communication-heavy graphs.

Recomputation must reproduce the original forward values. A changed random-number-generator state, mutable global, device move, or other side effect can make the recomputed path differ and can yield a silent gradient error. Framework checkpoint utilities preserve some random-number-generator state, but the exact device and state coverage is part of their documented contract (PyTorch Contributors 2026). Profile the real workload: checkpointing may lower peak memory while increasing arithmetic, communication exposure, or step time.

Eager, traced, and compiled are different axes

Early tensor systems exposed the representation they needed for optimization. Theano constructed symbolic graphs, and TensorFlow 1 executed a static dataflow graph through a session (Theano Development Team 2016; Abadi et al. 2016). Chainer and PyTorch made eager execution practical by recording the operations of each concrete run; PyTorch's design kept ordinary Python control flow while moving its tensor hot path below Python (Paszke et al. 2019). JAX took another route: transform pure functions through tracing, and compose differentiation, vectorization, and compilation as program transformations (Frostig et al. 2018).

This history is not a winner-takes-all framework war. Current systems combine immediate execution with one or more staged representations:

In PyTorch eager mode, autograd rebuilds a Function DAG on every recorded forward pass. torch.compile is a separate capture path: TorchDynamo extracts FX regions, records guards over assumptions such as type, shape, dtype, device, and stride, then sends captured graphs through AOTAutograd and a backend. With partial capture, unsupported code creates a graph break and execution can resume in another captured region. A failed guard can select another cached executable or trigger recompilation; requiring a full graph turns a break into an error (Ansel et al. 2024; PyTorch Contributors 2026).

JAX's grad and vmap are transformations, while jit stages and compiles a specialization. Compatible calls reuse the cached executable. Ordinary Python side effects occur during tracing rather than on every device execution, and a runtime data-dependent branch inside a jitted region needs a structured control flow primitive or a different staging boundary (JAX Authors 2026). TensorFlow 2 executes eagerly by default; GradientTape records eager TensorFlow operations, while tf.function traces and caches specialized TensorFlow graphs. XLA compilation is an additional choice, not a synonym for tf.function (TensorFlow Authors 2026; TensorFlow Authors 2026).

The useful comparison is therefore not “graph or tape.” It is which semantics are visible at trace time, what becomes part of the cache key, which effects remain observable, how failures fall back, and whether cold compile cost is amortized by warm execution.

What a framework must specify

A modern framework is a set of cross-cutting contracts rather than a neat stack of interchangeable layers:

1. A tensor representation defines shape, dtype, device, layout, strides, aliases, gradient state, and sometimes distributed placement.
2. An operator schema defines accepted inputs, broadcasting, dtype promotion, output metadata, mutation, and aliasing.
3. A dispatcher selects implementations not only by backend, but also for transformations and modes such as autograd, batching, autocast, functionalization, and tensor subclasses.
4. An autodiff transform supplies VJP and, when supported, JVP rules plus the saved-tensor contract.
5. A compiler needs abstract or fake-tensor behavior, graph lowering, shape reasoning, and a backend implementation.
6. A device runtime owns streams, events, synchronization, and a memory allocator. Asynchronous errors and allocator fragmentation surface through this layer rather than through the model equation.
7. A distributed tensor system attaches a logical layout to a device mesh and defines how operators propagate or change it.

A custom operator is where omissions become visible. Its registration may need an operator schema, backend kernel, fake or meta implementation for tracing, autograd rule, batching rule, autocast policy, and distributed partitioning rule. A registration checker can validate that these pieces are wired together; it cannot prove that the mathematical custom derivative is correct. That still requires a gradient check and representative eager-versus-compiled tests.

Distributed layout is tensor semantics

A distributed tensor represents one logical value over a device mesh. Its placement may be Shard, meaning each mesh coordinate stores part of a tensor dimension; Replicate, meaning each stores the complete value; or Partial, meaning each stores a contribution awaiting reduction. Partial is not a complete logical result.

Changing placement has communication semantics. Shard to Replicate generally requires an all-gather. Partial to Replicate requires an all-reduce, while Partial to Shard can use reduce-scatter. Changing the sharded dimension may require all-to-all. The framework can insert these collectives, but it cannot make their bytes free (PyTorch Contributors 2026). The physical topology and sustainable link rates from Chapter 62 remain constraints on the program.

GSPMD demonstrated how a compiler can propagate a small set of sharding annotations through a computation graph and emit a single-program, multiple-data execution (Xu et al. 2021). PyTorch's distributed tensors and FSDP integrate similar global-tensor semantics with the dispatcher, autograd, and allocator (Zhao et al. 2023). Manual per-device programs expose more control but also make collective order, global-versus-local shape, and loss normalization the programmer's responsibility. An inconsistent mesh or collective order can hang every rank even when each local tensor operation is valid. The parallelism algorithms built on these interfaces are developed in Chapter 10.

An operating checklist

Correctness comes before a speedup claim. For a new model path, operator, or framework backend:

• Compare eager and compiled outputs, gradients, state updates, and random number consumption on representative inputs.
• Run a directional finite difference and the framework's gradient check in double precision, away from nondifferentiable points.
• Cover shape, dtype, device, layout, noncontiguous views, broadcasting, zero-size tensors, aliasing, and mutation. Test NaN and infinity behavior deliberately.
• Measure cold compile and warm execution separately. Record cache keys or input signatures, graph-break locations, guard failures, and recompilation counts.
• Profile kernel time, host gaps, saved-tensor bytes, peak memory, allocator reserves, and communication volume. A faster kernel can still make the whole step slower.
• In distributed tests, assert global and local shapes, placement transitions, reduction denominators, mesh identity, collective order, checkpoint restore, and behavior after a rank or link failure.
• State the reproducibility target: tolerance-close, statistically equivalent, restart-equivalent, or bitwise on one fixed platform. Purity and fixed seeds help, but compiler paths, reduction order, world size, hardware, and releases can still change exact bits.

The serving incidents that arise from compilation caches, allocator behavior, and backend coverage continue in Chapter 31.

Further reading