How does JAX implement reverse-mode autodiff?

As forward mode plus a second pass over a smaller program. There is no separate reverse interpreter that walks your function backwards. Calling jax.vjp runs ad.linearize, which installs a LinearizeTrace holding two traces at once: the parent trace, where primal values are computed for real, and a separate tangent trace that is a jaxpr builder. Every primitive that arrives is split. Its primal half executes on the parent trace immediately; its linear half is recorded as equations on the tangent trace, along with the residuals it will need, which are the intermediate values from the forward pass. What comes back is the outputs plus a tangent jaxpr that is linear in the input tangents by construction. Only then does reverse mode happen, and it happens to that jaxpr rather than to your function: backward_pass3 transposes it. Because the tangent jaxpr is linear, transposing it is mechanical. jax.grad is a thin wrapper over jax.value_and_grad, which calls vjp and applies a cotangent of one, which is why grad insists the output is a scalar while the machinery underneath is perfectly happy with a vector.

What is the difference between a JVP rule and a transpose rule in JAX?

A JVP rule is per primitive and works on real values: it receives the primal arguments and their tangents and returns the primal output and the output tangent. It lives in the primitive_jvps dictionary. A transpose rule never sees your function at all. It only ever runs on the tangent jaxpr, and its job is to move a cotangent from an equation output back to that equation inputs. It lives in one of two dictionaries, primitive_transposes for the older style that receives UndefinedPrimal stand-ins and returns a list of cotangents, or fancy_transposes for the newer style that receives the accumulator objects themselves and adds into them in place. The consequence catches people out: only linear primitives need a transpose rule. sin has a JVP rule and no transpose rule anywhere in the JAX source tree, and differentiating it works fine, because linearizing sin turns it into a multiply by a residual and multiply is the primitive that gets transposed.

Why does JAX raise "Transpose rule for ... not implemented"?

Because a primitive reached the tangent jaxpr with no entry in the transpose registry. The message comes from one four-line function, get_primitive_transpose, which is a dictionary lookup wrapped in a try block that re-raises a KeyError as NotImplementedError with that text. The important word in it is transpose. A primitive with no derivative at all produces a different message, "Differentiation rule for ... not implemented", raised in the forward direction when the JVP trace fails to find the primitive in primitive_jvps. So the transpose message means the forward pass worked and produced a linear equation nobody taught JAX to run backwards, which is what you get from a custom primitive whose JVP rule emits the primitive itself linearly. There is a third message worth recognising: linearize also refuses when a primal output comes back unknown, and its text names the cause outright, an operation that does not support reverse-mode autodiff.

What is Zero in JAX autodiff and why does it exist?

It is a sentinel standing for a tangent that is identically zero, and it carries a type but no array. The class has one slot, the abstract value, and two module-level shorthands build one from a primal: p2tz for a tangent-space zero and p2cz for a cotangent-space zero. It exists so the machinery can skip work rather than compute with zeros. The JVP trace checks whether every input tangent is a Zero and, if so, rebinds the primitive on the parent trace and returns without consulting any rule at all, so a constant subgraph inside a differentiated function costs nothing. The tangent adder short-circuits both ways, returning the other operand when either side is a Zero, so a sum of tangents never binds an addition against a zero array. And the tracer constructor drops a Zero tangent entirely, returning the bare primal instead of a tracer. When a Zero has to become real, because some consumer needs an actual array, instantiate_zeros materialises it, and the instantiate argument on jvp decides which outputs get that treatment.

How do I register a derivative for a new JAX primitive?

Through one of six helpers, and which one you pick is a statement about the algebra of your primitive. defjvp takes one rule per positional argument, each returning that argument contribution to the output tangent, and sums them; that is the general case, and it is what sin and multiply both use. defjvp2 is the same but hands the rule the forward output as well, which saves recomputing it for something like tan. deflinear and deflinear2 are for a primitive linear in all its arguments: the JVP is just the primitive rebound on the tangents, so you supply only a transpose rule. defbilinear is for something linear in each argument separately but not jointly, matrix multiply and convolution being the two in the tree; you supply one transpose rule per argument and it registers them in both transpose registries. defjvp_zero declares that a primitive has no derivative, such as a comparison. There is also a seventh, lower-level path: registering directly in primitive_linearizations lets a primitive name its own residuals instead of having them discovered by tracing its JVP rule.

How JAX Builds a Backward Pass Out of a Forward One: ad.py, Line by Line

Forward-mode differentiation is easy to believe in: carry a derivative alongside every value and push both through the program together. Reverse mode, the one that makes training possible, looks like a different algorithm entirely. This 1162-line file says it is not. It runs the ordinary forward pass with one change, splitting each operation into a primal half that executes now and a linear half recorded as a second small program, and then it runs that second program backwards. Read whole at a pinned commit of jax-ml/jax, in the file order rather than the order things get called: the jvp entry point and its two-value tracer, the linearizer that builds the tangent jaxpr and prunes residuals it can forward instead of storing, the transpose interpreter that scans forward to find which equations touch a cotangent and then walks them in reverse, the three accumulator classes cotangents land in, the two registries a transpose rule can live in, the six helpers that install a derivative for a primitive, and the marker primitive whose transpose is your own custom_vjp function. The signature exhibit ports the real jvp and transpose rules for sin, mul and add and runs the whole pipeline on a function you pick: the primal equations, the tangent equations each rule emits, the linear jaxpr with its residuals named, then the reverse walk with every accumulation shown and the derivative it lands on checked against the analytic one. Closing on the counterfactual people get backwards: what actually happens when you differentiate a non-linear primitive that has no transpose rule at all.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of How JAX Builds a Backward Pass Out of a Forward One: ad.py, Line by Line — built to be understood, not skimmed.

Questions

How does JAX implement reverse-mode autodiff?: As forward mode plus a second pass over a smaller program. There is no separate reverse interpreter that walks your function backwards. Calling jax.vjp runs ad.linearize, which installs a LinearizeTrace holding two traces at once: the parent trace, where primal values are computed for real, and a separate tangent trace that is a jaxpr builder. Every primitive that arrives is split. Its primal half executes on the parent trace immediately; its linear half is recorded as equations on the tangent trace, along with the residuals it will need, which are the intermediate values from the forward pass. What comes back is the outputs plus a tangent jaxpr that is linear in the input tangents by construction. Only then does reverse mode happen, and it happens to that jaxpr rather than to your function: backward_pass3 transposes it. Because the tangent jaxpr is linear, transposing it is mechanical. jax.grad is a thin wrapper over jax.value_and_grad, which calls vjp and applies a cotangent of one, which is why grad insists the output is a scalar while the machinery underneath is perfectly happy with a vector.
What is the difference between a JVP rule and a transpose rule in JAX?: A JVP rule is per primitive and works on real values: it receives the primal arguments and their tangents and returns the primal output and the output tangent. It lives in the primitive_jvps dictionary. A transpose rule never sees your function at all. It only ever runs on the tangent jaxpr, and its job is to move a cotangent from an equation output back to that equation inputs. It lives in one of two dictionaries, primitive_transposes for the older style that receives UndefinedPrimal stand-ins and returns a list of cotangents, or fancy_transposes for the newer style that receives the accumulator objects themselves and adds into them in place. The consequence catches people out: only linear primitives need a transpose rule. sin has a JVP rule and no transpose rule anywhere in the JAX source tree, and differentiating it works fine, because linearizing sin turns it into a multiply by a residual and multiply is the primitive that gets transposed.
Why does JAX raise "Transpose rule for ... not implemented"?: Because a primitive reached the tangent jaxpr with no entry in the transpose registry. The message comes from one four-line function, get_primitive_transpose, which is a dictionary lookup wrapped in a try block that re-raises a KeyError as NotImplementedError with that text. The important word in it is transpose. A primitive with no derivative at all produces a different message, "Differentiation rule for ... not implemented", raised in the forward direction when the JVP trace fails to find the primitive in primitive_jvps. So the transpose message means the forward pass worked and produced a linear equation nobody taught JAX to run backwards, which is what you get from a custom primitive whose JVP rule emits the primitive itself linearly. There is a third message worth recognising: linearize also refuses when a primal output comes back unknown, and its text names the cause outright, an operation that does not support reverse-mode autodiff.
What is Zero in JAX autodiff and why does it exist?: It is a sentinel standing for a tangent that is identically zero, and it carries a type but no array. The class has one slot, the abstract value, and two module-level shorthands build one from a primal: p2tz for a tangent-space zero and p2cz for a cotangent-space zero. It exists so the machinery can skip work rather than compute with zeros. The JVP trace checks whether every input tangent is a Zero and, if so, rebinds the primitive on the parent trace and returns without consulting any rule at all, so a constant subgraph inside a differentiated function costs nothing. The tangent adder short-circuits both ways, returning the other operand when either side is a Zero, so a sum of tangents never binds an addition against a zero array. And the tracer constructor drops a Zero tangent entirely, returning the bare primal instead of a tracer. When a Zero has to become real, because some consumer needs an actual array, instantiate_zeros materialises it, and the instantiate argument on jvp decides which outputs get that treatment.
How do I register a derivative for a new JAX primitive?: Through one of six helpers, and which one you pick is a statement about the algebra of your primitive. defjvp takes one rule per positional argument, each returning that argument contribution to the output tangent, and sums them; that is the general case, and it is what sin and multiply both use. defjvp2 is the same but hands the rule the forward output as well, which saves recomputing it for something like tan. deflinear and deflinear2 are for a primitive linear in all its arguments: the JVP is just the primitive rebound on the tangents, so you supply only a transpose rule. defbilinear is for something linear in each argument separately but not jointly, matrix multiply and convolution being the two in the tree; you supply one transpose rule per argument and it registers them in both transpose registries. defjvp_zero declares that a primitive has no derivative, such as a comparison. There is also a seventh, lower-level path: registering directly in primitive_linearizations lets a primitive name its own residuals instead of having them discovered by tracing its JVP rule.

How JAX Builds a Backward Pass Out of a Forward One: ad.py, Line by Line

Questions

Related explainers