What is partial evaluation in JAX?

It is the mechanism that turns one traced function into two programs. Every value gets a single bit, known or unknown: known means the machine is holding the actual array right now, unknown means it has only a shape and a dtype. The rule that propagates the bit is stated in a comment above the code that implements it, in jax/_src/interpreters/partial_eval.py: "By default, if all the input tracers are known, then bind the primitive and consider all outputs known. Otherwise, stage the application into the jaxpr and consider all outputs unknown." So a known operation is actually executed during the trace rather than recorded, and one unknown input makes the whole application unknown, which makes its outputs unknown, which spreads downstream. The known operations run and produce arrays; the unknown ones become the equations of a second jaxpr. jax.grad, jax.checkpoint, lax.scan under differentiation, and a jit nested inside a grad are all built on that one rule.

What is a residual in JAX autodiff?

A residual is a value the known half computes and the unknown half reads. When the split rule meets an operation with a mix of known and unknown inputs, it instantiates every known operand: the value is pinned to a variable so an equation in the second program can name it, and the array is remembered on the side. The second jaxpr then takes those variables as its leading inputs and the first jaxpr grows matching extra outputs at the end of its output list. Under reverse-mode autodiff the known half is the forward pass and the unknown half is the linear program that consumes the cotangent, so the residuals are precisely the intermediate activations the backward pass needs, held in memory for the whole span between the two. That is why activation memory is a statement about how many residuals a cut produces, and why jax.checkpoint is not a separate feature but a predicate handed to the same splitter, deciding per equation whether its output is allowed to be a residual.

Why does JAX not constant-fold inside jit?

Because folding was moved to the compiler. Write x + jnp.sin(1.0) inside a jitted function and the jaxpr contains a sin equation applied to the literal 1.0, not the number 0.841. The staging trace calls try_constant_folding on every operation and its first line is a membership test against a table called const_fold_rules; a primitive with no entry gets an equation immediately. Across the whole of JAX at this commit there are three registrations: device_put, convert_element_type and stage. This is the visible face of omnistaging, the 2020 design change that made every operation inside jit stage regardless of whether its inputs were concrete. XLA folds the expression anyway, and keeping the tracer out of that business keeps it cheap on the hot path of every call. There is an opt-out flag, jax_eager_constant_folding, which short-circuits to eager evaluation whenever no operand is a tracer; it is off by default.

Where does a jaxpr get its variables?

A variable in a jaxpr is not a name. core.Var has a single slot, an abstract value, and carries no name, index, or definition site; its repr is Var(id=...), so two variables are distinguishable only because they are different Python objects, and the letters you see when a jaxpr prints are assigned by the printer as it walks. Under the partial-evaluation trace each tracer carries a recipe and there are four origins: a LambdaBinding becomes a fresh input variable, a ConstVar becomes a constvar deduplicated by the identity of the array and that is what a residual is, a FreeVar becomes an environment variable prepended to the inputs, and an equation recipe produces fresh output variables or a DropVar when nothing reads the output. A Literal produces no variable at all, because the atom in the equation is the value itself.

What is the difference between JaxprTrace and DynamicJaxprTrace?

They are the two interpreters defined in the JAX partial evaluation file, and they differ in whether they ever evaluate anything. JaxprTrace is the partial evaluator: each of its tracers carries a PartialVal that is either a real array or just an abstract value, and when an operation arrives it either performs it, if every input is known, or records it. It returns a jaxpr and a set of computed values. DynamicJaxprTrace is the staging interpreter jit installs, and it never performs anything: every operation becomes an equation, because at a jit boundary everything that crossed is a formal parameter and nothing is known. It also appends equations to a frame as it goes rather than hanging recipes on tracers and assembling the graph at the end, and it holds an effect-free equation by weak reference so an unused one is collected before the jaxpr is ever built.

Partial Evaluation: Splitting One Traced Call Into Two Programs

jax.grad hands you two programs out of one Python function: a forward pass that runs on the arguments you passed, and a linear pass that runs later when a cotangent arrives. Nobody wrote the second one. It comes out of partial evaluation, one file in jax/_src that gives every value a single bit, known or unknown, and propagates it with one rule: an operation is known only if every one of its inputs is. Known operations are executed during the trace; the rest become equations in a second jaxpr, together with the values they need from the first. Those crossing values are residuals, and they are exactly the activations a backward pass holds in memory, which is why rematerialisation is a predicate handed to the same splitter rather than a separate feature. The page reads the rule in the source, computes the known/unknown split on a small function you edit, and answers the rest of what the file settles: where a jaxpr gets its variables, why a constant is sometimes inlined and sometimes carried, and why the sine of a number you typed is still an equation inside jit.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Partial Evaluation: Splitting One Traced Call Into Two Programs — built to be understood, not skimmed.

Questions

What is partial evaluation in JAX?: It is the mechanism that turns one traced function into two programs. Every value gets a single bit, known or unknown: known means the machine is holding the actual array right now, unknown means it has only a shape and a dtype. The rule that propagates the bit is stated in a comment above the code that implements it, in jax/_src/interpreters/partial_eval.py: "By default, if all the input tracers are known, then bind the primitive and consider all outputs known. Otherwise, stage the application into the jaxpr and consider all outputs unknown." So a known operation is actually executed during the trace rather than recorded, and one unknown input makes the whole application unknown, which makes its outputs unknown, which spreads downstream. The known operations run and produce arrays; the unknown ones become the equations of a second jaxpr. jax.grad, jax.checkpoint, lax.scan under differentiation, and a jit nested inside a grad are all built on that one rule.
What is a residual in JAX autodiff?: A residual is a value the known half computes and the unknown half reads. When the split rule meets an operation with a mix of known and unknown inputs, it instantiates every known operand: the value is pinned to a variable so an equation in the second program can name it, and the array is remembered on the side. The second jaxpr then takes those variables as its leading inputs and the first jaxpr grows matching extra outputs at the end of its output list. Under reverse-mode autodiff the known half is the forward pass and the unknown half is the linear program that consumes the cotangent, so the residuals are precisely the intermediate activations the backward pass needs, held in memory for the whole span between the two. That is why activation memory is a statement about how many residuals a cut produces, and why jax.checkpoint is not a separate feature but a predicate handed to the same splitter, deciding per equation whether its output is allowed to be a residual.
Why does JAX not constant-fold inside jit?: Because folding was moved to the compiler. Write x + jnp.sin(1.0) inside a jitted function and the jaxpr contains a sin equation applied to the literal 1.0, not the number 0.841. The staging trace calls try_constant_folding on every operation and its first line is a membership test against a table called const_fold_rules; a primitive with no entry gets an equation immediately. Across the whole of JAX at this commit there are three registrations: device_put, convert_element_type and stage. This is the visible face of omnistaging, the 2020 design change that made every operation inside jit stage regardless of whether its inputs were concrete. XLA folds the expression anyway, and keeping the tracer out of that business keeps it cheap on the hot path of every call. There is an opt-out flag, jax_eager_constant_folding, which short-circuits to eager evaluation whenever no operand is a tracer; it is off by default.
Where does a jaxpr get its variables?: A variable in a jaxpr is not a name. core.Var has a single slot, an abstract value, and carries no name, index, or definition site; its repr is Var(id=...), so two variables are distinguishable only because they are different Python objects, and the letters you see when a jaxpr prints are assigned by the printer as it walks. Under the partial-evaluation trace each tracer carries a recipe and there are four origins: a LambdaBinding becomes a fresh input variable, a ConstVar becomes a constvar deduplicated by the identity of the array and that is what a residual is, a FreeVar becomes an environment variable prepended to the inputs, and an equation recipe produces fresh output variables or a DropVar when nothing reads the output. A Literal produces no variable at all, because the atom in the equation is the value itself.
What is the difference between JaxprTrace and DynamicJaxprTrace?: They are the two interpreters defined in the JAX partial evaluation file, and they differ in whether they ever evaluate anything. JaxprTrace is the partial evaluator: each of its tracers carries a PartialVal that is either a real array or just an abstract value, and when an operation arrives it either performs it, if every input is known, or records it. It returns a jaxpr and a set of computed values. DynamicJaxprTrace is the staging interpreter jit installs, and it never performs anything: every operation becomes an equation, because at a jit boundary everything that crossed is a formal parameter and nothing is known. It also appends equations to a frame as it goes rather than hanging recipes on tracers and assembling the graph at the end, and it holds an effect-free equation by weak reference so an unused one is collected before the jaxpr is ever built.

Partial Evaluation: Splitting One Traced Call Into Two Programs

Questions

Related explainers