std::autodiff: Differentiation in the Compiler

AI authored to showcase ironpad capabilities.

Introduction

Derivatives quietly run modern computing. Training a neural network is gradient descent on a loss function; a physics engine integrates forces, which are themselves derivatives of an energy; every "how does the output move when I nudge the input" question is asking for a derivative. So it matters, a great deal, how fast and how accurately you can compute them.

Rust's answer is unusually ambitious: put automatic differentiation inside the compiler. The std::autodiff feature (nightly, -Zautodiff) hands your function to Enzyme, an LLVM plugin that differentiates your code at the LLVM-IR level, after optimization. As of the Rust Project Goals update from April 2026, it is partially wired into CI and usable on nightly today if you supply the Enzyme library yourself. More on the real thing at the end.

The two algorithms the compiler automates are simple enough to build from scratch, and this notebook does exactly that. We will hand-roll forward-mode AD with dual numbers, plot a function against its own derivative without writing a line of calculus ourselves, quantify how exact it is against finite differences, then build the reverse-mode tape that powers every deep-learning framework in use today. Then we will look at what std::autodiff actually does, and why doing it in the compiler is the whole point.

No #[derive(Calculus)]. Just algebra that carries a derivative along for the ride.

Three ways to differentiate

Dual numbers make this look easy, so start with the two roads everyone tries first. AD only feels inevitable once you have watched both of them fail.

Symbolic differentiation is what a computer algebra system does: manipulate the formula. It is exact, and it swells. The derivative of a product of n factors has n terms; nest a few chain rules and the expression grows faster than you can simplify it. The deeper problem is that most code is not a formula at all. A function with a for loop, an early return, or a branch has no closed-form expression to manipulate; a CAS cannot even see it.

Numeric differentiation, the finite difference, treats the function as a black box: nudge the input, measure the slope. It works on any code, which is why everyone reaches for it, but its accuracy is capped by a tug-of-war we will quantify in a few cells: the best a perfectly tuned central difference can do in f64 is an absolute error around 10^{-11}, five orders of magnitude short of machine precision. And a gradient over n inputs costs n extra evaluations per point.

Automatic differentiation is the third road: differentiate the computation instead of the formula. Whatever its loops and branches, a program executes as a sequence of primitive operations (+, \times, \sin, \exp), and each primitive has a derivative you know in closed form. Chain them together as the program runs and you get derivatives exact to machine precision, at a small constant multiple of the original runtime, with no expression swell. Forward mode and reverse mode are simply the two directions you can apply the chain rule in: input to output, or output to input. That one choice of direction, we will see, is worth about eight orders of magnitude on a neural network.

Forward mode: differentiation with algebra

A dual number looks like a + b\,\varepsilon, where \varepsilon is a symbol with exactly one rule: \varepsilon^2 = 0. It is the same move as the complex numbers (i^2 = -1), pointed at calculus instead of geometry.

Watch what happens when you push a dual number through any analytic function and expand it as a Taylor series:

f(a + b\,\varepsilon) = f(a) + b\,f'(a)\,\varepsilon + \frac{b^2 f''(a)}{2}\,\varepsilon^2 + \cdots = f(a) + b\,f'(a)\,\varepsilon

Every term past the first derivative vanishes, because it carries an \varepsilon^2 or higher. So if you evaluate f at a + 1\cdot\varepsilon, the real part is f(a) and the \varepsilon coefficient is exactly f'(a). No limit, no small step, no symbolic manipulation: the derivative rides along in the \varepsilon slot, and the chain rule falls out of ordinary multiplication.

The whole engine is one struct and a handful of operator impls (see the Dual type in this notebook's shared source). Each operation encodes a derivative rule you already know by heart: the product rule lives in Mul, the quotient rule in Div, and \frac{d}{dx}\sin x = \cos x in Dual::sin. Our test function is a Gaussian-damped sine, a little wave packet:

f(x) = \sin(x)\,e^{-x^2/2}

Dual: forward mode in 40 lines (shared)⬡ shared

Watching the chain rule ride along

"The chain rule falls out of ordinary multiplication" is the kind of claim that deserves a slow-motion replay. Take the composition

s(x) = \sin(x^2)

at x = 1.2. Calculus says s'(x) = 2x\cos(x^2), and getting there by hand means consciously invoking the chain rule: derivative of the outside, times derivative of the inside.

Dual numbers never invoke it. Each operation applies only its own local rule: Mul applies the product rule to its two operands and nothing else, sin multiplies the incoming derivative by \cos of the incoming value and nothing else. No operation knows what expression it is embedded in. The composition of the local rules is the chain rule; it happens by construction, in the order the program already executes.

That locality is the load-bearing property. It is why AD scales to functions with loops and branches: the tape of executed primitives is always a straight line, whatever the control flow that produced it. The next cell traces every intermediate of \sin(x^2) so you can watch the derivative accumulate one primitive at a time.

Tracing sin(x^2), one primitive at a time
f and f' at a point
f vs f' across the packet

How exact is "exact"? Dual numbers vs finite differences

Before dual numbers, the default way to get a derivative you did not want to work out by hand was the finite difference: nudge the input by a small h and measure the slope. The central difference is the good version,

f'(x) \approx \frac{f(x+h) - f(x-h)}{2h}

and it looks unbeatable until you actually chase the accuracy down. It is caught between two errors pulling in opposite directions. Truncation error comes from the formula being an approximation; a Taylor expansion pins it at \frac{h^2}{6}f'''(\xi), so it shrinks like O(h^2) as you make h smaller. Roundoff error comes from computing f(x+h) - f(x-h) in floating point: you are subtracting two nearly equal numbers, which is catastrophic cancellation, then dividing by a tiny 2h, which inflates the surviving machine noise to O(\varepsilon_{\text{mach}}/h) and grows as h shrinks.

Add the two and you get a valley. Minimizing C_1 h^2 + C_2\,\varepsilon_{\text{mach}}/h gives an optimal step and a hard floor:

h_\star \propto \varepsilon_{\text{mach}}^{1/3} \approx 6\times 10^{-6}, \qquad \text{best error} \propto \varepsilon_{\text{mach}}^{2/3} \approx 4\times 10^{-11}

So the best a central difference can do, tuned perfectly, is an absolute error around 10^{-11}. The dual number carries no step size at all, so it has no truncation error and nothing to cancel: it should land near machine precision, roughly 10^{-16}. The next cell puts all three side by side and lets the numbers make the argument.

Accuracy table + error valley

Forward mode's party trick: derivatives along a direction

Every seed so far has been 1.0 on a single input. Nothing says it has to be. Seed the inputs of a multivariate function with an arbitrary vector v = (v_x, v_y, \ldots) and linearity does the rest: the \varepsilon slot of the output comes back holding

\nabla f \cdot v

the directional derivative along v, in one pass, no matter how many inputs the function has. In AD jargon this is a Jacobian-vector product (JVP), and it is forward mode's native operation. The full gradient is just the special case where you ask for n directional derivatives along the n basis vectors, which is exactly why it costs n passes.

For a mental model: forward mode answers the question "if I wiggle the inputs in this particular direction, how does the output move?" One question, one pass, no memory beyond the values themselves. Plenty of real work only ever needs that one question: a line search asking "does the loss improve along this step?", a sensitivity analysis along one trajectory, the inner loop of a Hessian-vector product. For those, forward mode is the right tool rather than a slow fallback.

The transposed question, "which input wiggles affect this output?", is a different beast entirely, and answering it efficiently is reverse mode's whole reason to exist. First, the party trick in action.

One pass, one direction

Reverse mode, and why your GPU only does it one way

Forward mode has a scaling problem hiding in plain sight. Each pass seeds one input's \varepsilon to 1 and reads back one input's partial derivative. A function with n inputs needs n passes to get the full gradient. For our little g(x,y,z) that is 3 passes, no big deal. For a neural network with 10^8 parameters, that is 10^8 forward passes per gradient, and training would simply never finish.

Reverse mode flips the cost. It runs one forward pass to compute the value, recording every operation onto a tape (a Wengert list), then one backward pass that walks the tape in reverse, seeds the output's sensitivity to 1, and accumulates an adjoint \bar{v} = \partial(\text{output})/\partial v into every intermediate. One forward plus one backward yields the entire gradient, no matter how many inputs there are. That single asymmetry, cheap gradients of one scalar with respect to millions of parameters, is the whole reason backpropagation exists and the reason every deep-learning framework is reverse-mode under the hood.

The mechanics are humble. Each node remembers its parents and the local partial \partial(\text{this})/\partial(\text{parent}); the backward sweep multiplies adjoints along the edges and sums where paths merge, which is just the multivariate chain rule in a for loop. The shared Tape and Var types do exactly this in about sixty lines. The next cell differentiates

g(x, y, z) = x\,y + \sin(x\,z) + y^2

and pulls all three partials out of a single backward pass, cross-checking them against forward mode and against the gradient worked out by hand.

Tape: reverse mode on a Wengert list (shared)⬡ shared
One backward pass, every partial

Anatomy of a backward sweep

That table came out of one backward pass, whose mechanics are simpler than its reputation. For g(x,y,z) = xy + \sin(xz) + y^2, the forward evaluation appends nine nodes to the tape, in execution order: the three leaves x, y, z, then xy, xz, \sin(xz), the first sum, y^2, and the final sum. Each non-leaf node records exactly two things per parent: who it came from, and the local partial \partial(\text{this})/\partial(\text{parent}) frozen at the values from the forward pass.

The backward sweep is bookkeeping. Seed the output's adjoint \bar g = 1, walk the list in reverse, and at each node hand its adjoint down to each parent, multiplied by the local partial on that edge. Where two paths merge, the contributions add: y feeds both xy and y^2, so its adjoint arrives in two installments, \bar y = x \cdot 1 + 2y \cdot 1. That sum-over-paths is the multivariate chain rule, executed as a for loop over a Vec.

For a mental model, think of an adjoint as an exchange rate: \bar v says how many units the output moves per unit wiggle of node v. The sweep starts with the output at 1 (it moves one-for-one with itself) and propagates the rate backwards through every edge, discounting by each edge's local partial.

Two things fall out of this for free. First, the sweep prices every node, not just the leaves: intermediates get adjoints too, and reading them is how debuggers and interpretability tools answer "which part of the computation mattered?". Second, the sweep touches each tape entry exactly once, which is where reverse mode's fixed cost comes from. The next cell dumps the whole priced tape.

The priced tape: adjoints at every node

The cost model

Time to make the "forward has a scaling problem" claim precise. For a function f: \mathbb{R}^n \to \mathbb{R}^m:

  • Forward mode costs one pass per input you care about: the full Jacobian takes n passes, each a small constant multiple (roughly 2-3x) of evaluating f once. Memory overhead: essentially zero. Duals live and die in registers.
  • Reverse mode costs one pass per output: the full Jacobian takes m sweeps over a tape built by one forward evaluation. Each sweep is again a small constant multiple of an evaluation. Memory overhead: the tape records every primitive executed, so it grows with runtime, not with n.
shape of fwinnerwhy
n small (a handful of inputs)forwardno tape, no memory, embarrassingly simple
m = 1, n huge (a loss function)reversethe entire gradient for ~3 evaluations' work
n \approx m, both largeneither dominatespick by memory budget, or mix modes
Hessian-vector productsforward over reversedifferentiate the gradient along a direction

The bottom-left cell of that table is deep learning. A loss over 10^8 parameters is f: \mathbb{R}^{10^8} \to \mathbb{R}: forward mode needs 10^8 passes, reverse needs one. That ratio, n versus a constant, is the entire economic basis of backpropagation, and it is why every framework is reverse-mode at its core.

The price reverse pays is that tape. Its memory scales with the number of operations executed, which for a big model means activations for every layer held live until the backward sweep consumes them. The standard escape is checkpointing: store the tape only at intervals and recompute the segments between checkpoints during the sweep, trading roughly one extra forward pass for a memory footprint that drops by the checkpoint spacing. Every serious framework does this; so do the HPC codes Enzyme grew up in.

The next cell makes the n-versus-constant asymmetry concrete with primitive-operation counts on our own Dual and Tape.

Counting the work: n passes vs one sweep

Choosing a mode in practice

The theory compresses to three working rules.

Rule 1: gradients of a scalar go reverse. Loss functions, energies, log-likelihoods: anything shaped \mathbb{R}^n \to \mathbb{R} with n bigger than a dozen. This is the overwhelmingly common case, which is why "autodiff" and "backprop" read as synonyms from inside machine learning.

Rule 2: few inputs, or one direction, go forward. Sensitivities of a simulation with three tunable parameters, a line search along one step direction, anything where you would only ask for a handful of the n basis passes anyway. You skip the tape, the memory, and a surprising amount of implementation complexity; our Dual is 40 lines against the tape's 60, and the gap widens fast in production systems.

Rule 3: second-order information mixes them. The Hessian-vector product Hv, the workhorse of second-order optimizers and curvature analysis, is forward mode applied over reverse mode: differentiate the gradient function along direction v. One forward-over-reverse pass costs a small multiple of one gradient, versus n gradient evaluations to materialize H column by column. Composability of the two modes is a primitive.

This is also exactly the shape of the modern framework APIs. JAX exposes jvp (forward) and vjp (reverse) as its two primitives and builds grad, jacfwd, jacrev, and hessian as compositions of them. PyTorch defaults to reverse and grew torch.func.jvp for the forward cases. The vocabulary you built in this notebook, seeds, tapes, adjoints, is their vocabulary, minus the marketing.

So: two algorithms, one choice of direction, clear rules for which to reach for. Now for what the Rust compiler does with them.

What std::autodiff actually does

Everything above we wrote by hand. The compiler feature runs the same two algorithms, except Enzyme differentiates the LLVM IR after the optimizer has already run, and that ordering is the entire point. Our Dual type, and every source-level autodiff library, differentiates the source and then hopes the optimizer cleans up the result. Enzyme differentiates code that has already been inlined, constant-folded, and vectorized, so the derivative composes with the optimizations instead of fighting them. That is how it stays fast enough to be worth it.

You opt in per function with an attribute that names the generated derivative and gives an activity type to each argument and return value:

#![feature(autodiff)]
use std::autodiff::autodiff_reverse;

// Enzyme synthesizes `d_square` from `square`, at the LLVM-IR level.
#[autodiff_reverse(d_square, Active, Active)]
fn square(x: f64) -> f64 { x * x }
// The reverse pass gives df/dx; at x = 3 that is 6.0.

The activity types are the vocabulary for telling Enzyme what to do with each slot:

  • Dual carries a forward-mode derivative alongside the value, exactly our \varepsilon slot.
  • Active marks a scalar that is differentiated in reverse mode and returned by value.
  • Duplicated pairs an input with a shadow buffer that receives its gradient, the reverse-mode story for arrays and structs.
  • Const says do not differentiate this one.

Can this run in a cell? Yes, with real compiler support: ironpad pins a dedicated nightly toolchain that carries Enzyme, injects the crate-root #![feature(autodiff)] gate, and builds autodiff cells with fat LTO. The companion cannon notebook does exactly that, differentiating a timestepped simulator live in your browser. The cells in this notebook stay hand-rolled on purpose: building the two algorithms yourself is how you feel the way they work.

The roadmap is ambitious. The same effort is landing std::offload (mark a function to run on the GPU through LLVM's offload backend) and a batching attribute (fuse N scalar calls into one vectorized call). Compose #[offload] with #[autodiff] and you get a differentiated function running on a GPU, which is to say Rust-native machine learning and scientific computing with the gradients handled by the compiler instead of a Python framework.

The signatures Enzyme writes for you

The activity annotations are easier to read once you see the functions they generate. Simplified (the exact ABI has more cases), the three shapes you will actually use:

Forward mode, Dual: the generated function threads a tangent alongside every value, which is precisely our Dual::var and .eps in compiler-generated form.

#[autodiff_forward(d_f, Dual, Dual)]
fn f(x: f64) -> f64 { x.sin() * x }

// generated, morally:
// fn d_f(x: f64, dx: f64) -> (f64, f64)
//        value ─┘    └─ seed        └─ (f(x), f'(x)·dx)

Call d_f(x, 1.0) and the second slot of the result is f'(x): the seed is our \varepsilon coefficient, and passing a non-unit seed gives you the JVP from a few cells back.

Reverse mode on scalars, Active: the generated function takes an adjoint seed for the output and returns the input's adjoint alongside the value.

#[autodiff_reverse(d_f, Active, Active)]
fn f(x: f64) -> f64 { x.sin() * x }

// generated, morally:
// fn d_f(x: f64, seed: f64) -> (f64, f64)   // (f(x), seed · f'(x))

The seed is \bar{g}, the value our tape initialized to 1.0 before the sweep.

Reverse mode on buffers, Duplicated: the shape that matters for real workloads. The input is paired with a caller-owned shadow buffer, and the backward pass accumulates the gradient into it, no tape object in sight because Enzyme synthesized the sweep as straight-line code.

#[autodiff_reverse(d_loss, Duplicated, Active)]
fn loss(w: &[f64]) -> f64 { /* ... */ }

// generated, morally:
// fn d_loss(w: &[f64], dw: &mut [f64], seed: f64) -> f64
//                      └─ gradient lands here, accumulated

That dw buffer is our grad vector with the indices designed away. Everything this notebook hand-rolled maps one-to-one onto these signatures; the compiler's version simply exists after LLVM has optimized the primal, which is the performance story of the previous section.

Where it stands, and who to thank

The status, from the Rust Project Goals update of April 2026 (published 2026-05-18):

  • Partially in CI. Autodiff is enabled on some Linux and MinGW runners, not yet everywhere.
  • Usable today, if you bring your own Enzyme. You can download a libEnzyme artifact, drop it into your nightly toolchain, and build with -Zautodiff=Enable and lto = "fat". It works; it is just not yet the one-command experience.
  • macOS is the blocker. A turnkey rollout is waiting on an LLVM static-to-dynamic linking change on macOS. Once that lands, a rustup component will auto-download Enzyme and the whole thing collapses to "just works on nightly."
  • Rough edges remain: no dyn Trait arguments, fat LTO required, debug builds are flakier, and a few impl-block miscompiles are still open.

Nearly all of this is the work of Manuel Drehwald (@ZuseZ4), roughly five days a week, sponsored by LLNL and the University of Toronto, after about three years spent building the LLVM autodiff tooling and a compiler fork before any of it went upstream. If you want the full tour from the source, his RustWeek 2026 talk, "std::autodiff: computing derivatives with your compiler", is the one to watch.

The receipts, if you want to follow along or contribute:

You just built, by hand, the two algorithms a compiler is learning to generate for you. When it ships, you will know exactly what it is doing under the hood, because you already did it once yourself.