Training is gradient descent, and gradient descent needs gradients, so every ML framework has an automatic differentiation engine at its core. What is less obvious is that the representation a framework chooses for autodiff is one of its most consequential design decisions. It decides what control flow you may write, how much the compiler can optimize, where the engine sits in the stack, and ultimately what code you are allowed to differentiate at all. PyTorch, JAX, and compiler-level tools like Enzyme make three different choices, and the differences explain a lot of their downstream behavior.
Forward and reverse mode
Automatic differentiation applies the chain rule mechanically to a program. There are two directions to walk the chain, and the choice is purely about cost.
Forward mode propagates a derivative alongside each value, from inputs to outputs. Its cost scales with the number of inputs: one pass per input direction. Reverse mode computes the output first, then propagates sensitivities backward from outputs to inputs. Its cost scales with the number of outputs.
Concretely, write the computation as y = f(θ). Forward mode seeds an input direction and pushes it through, accumulating a Jacobian-vector product (JVP). Reverse mode seeds the output and pulls a vector-Jacobian product (VJP) back to the inputs.
For training, the loss is a scalar, so there is exactly one output and millions of parameters. Reverse mode is the only sane choice: one backward pass computes the gradient of the loss with respect to every parameter. This is why every ML framework is reverse-mode at heart.
The three representations
1. The runtime tape (PyTorch). PyTorch's autograd records operations into a tape as they execute, building a graph of tensors and their grad-fns. The tape is built at runtime, which means control flow is unrestricted: whatever Python does, the tape records it. The cost is that the tape is a runtime data structure, and the graph is only as good as what was recorded. The tape approach is why PyTorch supports arbitrary Python control flow in the forward pass.
2. The functional trace-and-transform (JAX). JAX traces the function into a jaxpr, then transforms the jaxpr into its reverse-mode form. Because the trace happens ahead of time, the compiler sees the whole program and can optimize the backward pass alongside the forward. The cost is the trace boundary: data-dependent control flow must be expressed through jax.lax primitives, and the refusal to trace it is the TracerBoolConversionError.
3. Source (IR) transformation (Enzyme). Enzyme differentiates at the LLVM IR level, transforming the machine code itself. It needs no tape and no trace; it works on whatever the compiler produced. The cost is that it must handle the full complexity of real IR: memory aliasing, control flow, and the whole instruction set.
Autodiff representations are like three ways of remembering a journey. PyTorch records a video as you drive (tape). JAX draws a map first, then plans the reverse route on the map (trace-and-transform). Enzyme takes the finished road and computes the reverse journey from the road signs themselves (IR transform).
The representation a framework chooses for autodiff decides what control flow you may write, how much the compiler can optimize, and ultimately what code you are allowed to differentiate at all.