Does the order of jit, grad, and vmap change the result in JAX?

For grad and vmap, yes — reordering them changes not just the number but what is even being computed. jax.grad "Creates a function that evaluates the gradient of fun," and the function it wraps "should return a scalar." vmap "Creates a function which maps fun over argument axes," which turns a scalar-returning function into a vector-returning one. So vmap(grad(f)) differentiates first, one example at a time, then batches those independent scalar gradients into an array — grad never sees anything but a scalar. grad(vmap(f)) tries to differentiate the batched, vector-valued function directly, which grad refuses outright, so the composition needs a reduction (a sum or mean) added before grad can even run. jit is different: it is value-transparent, so its position among grad and vmap does not change the answer, only when the computation gets compiled.

Why does grad(vmap(f)) fail in JAX unless you add a reduction?

Because vmap turns a per-example scalar loss into a vector of per-example losses — one number per row of the batch — and grad only ever differentiates a function that "should return a scalar." Feed grad something that returns an array instead and there is no single slope to compute, so it stops you before you can run it. The fix everyone reaches for is to wrap the vmapped function in a sum or a mean before differentiating: grad(lambda w, xs, ys: vmap(loss)(w, xs, ys).sum())(w, xs, ys). That composition works, but notice the cost: you had to choose sum vs mean, and that choice changes the returned number by exactly the batch size, a decision vmap(grad(f)) never forces on you because it never aggregates.

What is the difference between vmap(grad(f)) and grad(vmap(f))?

vmap(grad(f)) computes grad(f) once per example and stacks the results — the standard idiom for per-example gradients, returning one gradient value for every row in the batch. grad(vmap(f)) needs vmap(f)'s vector output collapsed to a scalar first, so it returns a single aggregate gradient — by linearity, exactly the sum (or mean) of the per-example gradients you would have gotten from vmap(grad(f)), just added together into one number instead of kept apart. Same underlying math, completely different shape: an array of N sensitivities versus one combined sensitivity — and the second shape only exists because a reduction was forced into the pipeline.

What is a JVPTracer and a BatchTracer in JAX?

They are how grad and vmap actually run your Python function without you rewriting it. Call grad(f) or vmap(f) and JAX re-executes f once, but with ordinary arrays swapped for tracer objects: grad’s JVPTracer (defined in jax/_src/interpreters/ad.py) carries a __slots__ pair of primal and tangent — the value and its derivative, riding together through every operation. vmap’s BatchTracer (jax/_src/interpreters/batching.py) carries val and batch_dim — the value and which axis, if any, is the mapped batch dimension. Composing transforms means nesting these tracers: under vmap(grad(f)), each BatchTracer's val is itself a JVPTracer; under grad(vmap(f)), it is the other way around — and that nesting order is exactly why the two compositions compute different things.

When should you use jacfwd vs jacrev in JAX?

It depends on the shape of the Jacobian, not on taste. Both are built by composing vmap with one direction of automatic differentiation — jacfwd vmaps jvp (forward-mode) and jacrev vmaps vjp (reverse-mode) — and JAX's own docs are direct about the trade-off: "jacfwd uses forward-mode automatic differentiation, which is more efficient for 'tall' Jacobian matrices (more outputs than inputs), while jacrev uses reverse-mode, which is more efficient for 'wide' Jacobian matrices (more inputs than outputs)." Forward-mode sweeps one direction per input, so its cost scales with the number of inputs; reverse-mode sweeps one direction per output, so its cost scales with the number of outputs — pick whichever dimension is smaller.

jit(grad(vmap(f))): Why Transform Order Changes the Answer

jit, grad, and vmap all take a function and return a function, which makes them feel like decorators you can stack in any order. They are not: vmap(grad(f)) computes one gradient per example — the standard trick behind per-example gradients — while grad(vmap(f)) does not even type-check until you bolt on a reduction, and which reduction you pick changes the number you get. We ground it in JAX’s own tracer classes, reorder a live jit/grad/vmap chip stack and watch the trace-stack diagram and a real numeric example recompute together, then find the same order-sensitivity again in jacfwd vs jacrev.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of jit(grad(vmap(f))): Why Transform Order Changes the Answer — built to be understood, not skimmed.

Questions

Does the order of jit, grad, and vmap change the result in JAX?: For grad and vmap, yes — reordering them changes not just the number but what is even being computed. jax.grad "Creates a function that evaluates the gradient of fun," and the function it wraps "should return a scalar." vmap "Creates a function which maps fun over argument axes," which turns a scalar-returning function into a vector-returning one. So vmap(grad(f)) differentiates first, one example at a time, then batches those independent scalar gradients into an array — grad never sees anything but a scalar. grad(vmap(f)) tries to differentiate the batched, vector-valued function directly, which grad refuses outright, so the composition needs a reduction (a sum or mean) added before grad can even run. jit is different: it is value-transparent, so its position among grad and vmap does not change the answer, only when the computation gets compiled.
Why does grad(vmap(f)) fail in JAX unless you add a reduction?: Because vmap turns a per-example scalar loss into a vector of per-example losses — one number per row of the batch — and grad only ever differentiates a function that "should return a scalar." Feed grad something that returns an array instead and there is no single slope to compute, so it stops you before you can run it. The fix everyone reaches for is to wrap the vmapped function in a sum or a mean before differentiating: grad(lambda w, xs, ys: vmap(loss)(w, xs, ys).sum())(w, xs, ys). That composition works, but notice the cost: you had to choose sum vs mean, and that choice changes the returned number by exactly the batch size, a decision vmap(grad(f)) never forces on you because it never aggregates.
What is the difference between vmap(grad(f)) and grad(vmap(f))?: vmap(grad(f)) computes grad(f) once per example and stacks the results — the standard idiom for per-example gradients, returning one gradient value for every row in the batch. grad(vmap(f)) needs vmap(f)'s vector output collapsed to a scalar first, so it returns a single aggregate gradient — by linearity, exactly the sum (or mean) of the per-example gradients you would have gotten from vmap(grad(f)), just added together into one number instead of kept apart. Same underlying math, completely different shape: an array of N sensitivities versus one combined sensitivity — and the second shape only exists because a reduction was forced into the pipeline.
What is a JVPTracer and a BatchTracer in JAX?: They are how grad and vmap actually run your Python function without you rewriting it. Call grad(f) or vmap(f) and JAX re-executes f once, but with ordinary arrays swapped for tracer objects: grad’s JVPTracer (defined in jax/_src/interpreters/ad.py) carries a __slots__ pair of primal and tangent — the value and its derivative, riding together through every operation. vmap’s BatchTracer (jax/_src/interpreters/batching.py) carries val and batch_dim — the value and which axis, if any, is the mapped batch dimension. Composing transforms means nesting these tracers: under vmap(grad(f)), each BatchTracer's val is itself a JVPTracer; under grad(vmap(f)), it is the other way around — and that nesting order is exactly why the two compositions compute different things.
When should you use jacfwd vs jacrev in JAX?: It depends on the shape of the Jacobian, not on taste. Both are built by composing vmap with one direction of automatic differentiation — jacfwd vmaps jvp (forward-mode) and jacrev vmaps vjp (reverse-mode) — and JAX's own docs are direct about the trade-off: "jacfwd uses forward-mode automatic differentiation, which is more efficient for 'tall' Jacobian matrices (more outputs than inputs), while jacrev uses reverse-mode, which is more efficient for 'wide' Jacobian matrices (more inputs than outputs)." Forward-mode sweeps one direction per input, so its cost scales with the number of inputs; reverse-mode sweeps one direction per output, so its cost scales with the number of outputs — pick whichever dimension is smaller.

jit(grad(vmap(f))): Why Transform Order Changes the Answer

Questions

Related explainers