How does jax.vmap work internally?

It installs an interpreter. vmap wraps your function so that, while it runs, the current trace is a BatchTrace, and every argument you asked to map becomes a BatchTracer: a pair of the real array and one integer, the batch dimension, saying which axis of that array is the mapped one. The tracer advertises the per-example shape rather than the real one, so your code sees the unbatched program it was written as. Every time a primitive is bound, BatchTrace.process_primitive unzips those pairs into values and dimensions, looks the primitive up in a dictionary called fancy_primitive_batchers, and calls the rule it finds there with the signature (axis_data, vals, dims, **params), getting back the output value and the output batch dimension. The result is wrapped in a fresh BatchTracer and the next primitive repeats it. Nothing is vectorised in a loop and nothing is unrolled; the whole transformation is a table lookup per operation plus some axis bookkeeping at the boundaries.

What does "Batching rule for X not implemented" mean in JAX?

It means a primitive with no entry in the batching table received an argument that carries a batch dimension. The code that raises it is the last branch of BatchTrace.process_primitive, and reading the two branches above it is what makes the error make sense. First, if the primitive is in fancy_primitive_batchers, its rule runs. Second, if every input dimension is None, which is the case the comment on that branch calls "Not all primitives have batching rules defined", the primitive is re-bound on the parent trace untouched, because there is nothing to batch. Only when both fail does the NotImplementedError fire. So a rule-less primitive is perfectly usable under vmap; it breaks the moment a mapped value reaches it. The fix is either to keep the batched value away from it, or to register a rule.

What does in_axes=None mean in jax.vmap?

It means that argument is not mapped, and it is handled before any tracing happens. The function to_elt turns each argument into what the trace will see: when the spec is an integer, it is canonicalised against the argument rank first, so a negative axis is allowed, and the argument becomes a BatchTracer carrying that axis. When the spec is None, the argument is returned exactly as it came in, an ordinary array rather than a tracer. That is why an unmapped argument costs nothing: it never enters the batching machinery, and every rule that meets it sees a dimension of None and takes its unbatched path.

Why does vmap move my batch axis to the front?

Because one of the rules could not keep it where it was. The rule behind every elementwise binary operation checks whether the batched arguments already agree: same shape, same batch dimension, with scalars ignored. If they agree it calls the primitive unchanged and reports the same dimension back, so the axis stays put. If they disagree, which is what happens when you map two arguments on different axes or map one and not the other, it moves every non-scalar argument so the batch is axis 0, inserting a size-one axis for the unmapped ones and letting the primitive broadcast, and reports 0. From then on the batch axis is at the front. The move itself is one call to moveaxis in a four-line helper named bdim_at_front, and the transposes it introduces are the reason a vmapped program can be slower than you expected. Nothing about the final answer changes: the very last thing vmap does is match each output axis to your out_axes.

How do I write a batching rule for a custom JAX primitive?

Three ways, in ascending order of effort. If your primitive is elementwise and unary, one line of defvectorized registers a rule that asserts the input dimensions all agree and passes the dimension straight through. If it is n-ary with broadcasting, defbroadcasting gives you the agree-or-move-to-front behaviour described above. If it reduces over axes, defreducer shifts every reduced axis past the batch dimension and recomputes where the batch lands among the axes that survive. Anything else gets an entry in the table by hand, written as (axis_data, vals, dims, **params) returning a value and an output dimension. The older spelling, an assignment to primitive_batchers, still works: it is a proxy object whose setter wraps your rule so that the all-unmapped case is handled for you and axis_data is dropped before your rule sees it.

Where vmap Puts the Axis: batching.py, Line by Line

jax.vmap adds a dimension to your arrays and then spends the rest of the call deciding where that dimension should sit. This is the 789-line file that decides. Every value under a vmap is a pair, the array and one integer saying which of its axes is the batch, and every primitive gets a rule that answers one question: given where the batch axis is on the inputs, where is it on the output? Read the file and vmap stops being a vectoriser and becomes a table you can extend.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Where vmap Puts the Axis: batching.py, Line by Line — built to be understood, not skimmed.

Questions

How does jax.vmap work internally?: It installs an interpreter. vmap wraps your function so that, while it runs, the current trace is a BatchTrace, and every argument you asked to map becomes a BatchTracer: a pair of the real array and one integer, the batch dimension, saying which axis of that array is the mapped one. The tracer advertises the per-example shape rather than the real one, so your code sees the unbatched program it was written as. Every time a primitive is bound, BatchTrace.process_primitive unzips those pairs into values and dimensions, looks the primitive up in a dictionary called fancy_primitive_batchers, and calls the rule it finds there with the signature (axis_data, vals, dims, **params), getting back the output value and the output batch dimension. The result is wrapped in a fresh BatchTracer and the next primitive repeats it. Nothing is vectorised in a loop and nothing is unrolled; the whole transformation is a table lookup per operation plus some axis bookkeeping at the boundaries.
What does "Batching rule for X not implemented" mean in JAX?: It means a primitive with no entry in the batching table received an argument that carries a batch dimension. The code that raises it is the last branch of BatchTrace.process_primitive, and reading the two branches above it is what makes the error make sense. First, if the primitive is in fancy_primitive_batchers, its rule runs. Second, if every input dimension is None, which is the case the comment on that branch calls "Not all primitives have batching rules defined", the primitive is re-bound on the parent trace untouched, because there is nothing to batch. Only when both fail does the NotImplementedError fire. So a rule-less primitive is perfectly usable under vmap; it breaks the moment a mapped value reaches it. The fix is either to keep the batched value away from it, or to register a rule.
What does in_axes=None mean in jax.vmap?: It means that argument is not mapped, and it is handled before any tracing happens. The function to_elt turns each argument into what the trace will see: when the spec is an integer, it is canonicalised against the argument rank first, so a negative axis is allowed, and the argument becomes a BatchTracer carrying that axis. When the spec is None, the argument is returned exactly as it came in, an ordinary array rather than a tracer. That is why an unmapped argument costs nothing: it never enters the batching machinery, and every rule that meets it sees a dimension of None and takes its unbatched path.
Why does vmap move my batch axis to the front?: Because one of the rules could not keep it where it was. The rule behind every elementwise binary operation checks whether the batched arguments already agree: same shape, same batch dimension, with scalars ignored. If they agree it calls the primitive unchanged and reports the same dimension back, so the axis stays put. If they disagree, which is what happens when you map two arguments on different axes or map one and not the other, it moves every non-scalar argument so the batch is axis 0, inserting a size-one axis for the unmapped ones and letting the primitive broadcast, and reports 0. From then on the batch axis is at the front. The move itself is one call to moveaxis in a four-line helper named bdim_at_front, and the transposes it introduces are the reason a vmapped program can be slower than you expected. Nothing about the final answer changes: the very last thing vmap does is match each output axis to your out_axes.
How do I write a batching rule for a custom JAX primitive?: Three ways, in ascending order of effort. If your primitive is elementwise and unary, one line of defvectorized registers a rule that asserts the input dimensions all agree and passes the dimension straight through. If it is n-ary with broadcasting, defbroadcasting gives you the agree-or-move-to-front behaviour described above. If it reduces over axes, defreducer shifts every reduced axis past the batch dimension and recomputes where the batch lands among the axes that survive. Anything else gets an entry in the table by hand, written as (axis_data, vals, dims, **params) returning a value and an output dimension. The older spelling, an assignment to primitive_batchers, still works: it is a proxy object whose setter wraps your rule so that the all-unmapped case is handled for you and axis_data is dropped before your rule sees it.

Where vmap Puts the Axis: batching.py, Line by Line

Questions

Related explainers