The rule per primitive

Print the type of an argument from inside a vmapped function and a BatchTracer comes back with three slots on it: the real value, an integer naming which of that value's axes is the batch, and a source location kept for error messages. The array underneath is the entire batch. Only the integer says which axis it is.

The tracer then hides that axis from you. Map axis 1 of a (3, 7, 5) array and the tracer reports its shape as (3, 5) while the value it holds is still (3, 7, 5) with batch_dim 1. You wrote the function for one example, and one example is the shape it gets handed, which is why the body needs no rewriting.

The second argument in the run below is the more interesting one. Give it None in in_axes and it does not arrive as a tracer at all; it arrives as an ArrayImpl, the concrete array, unwrapped. An unbatched argument is not a batch of size one and not a broadcast copy. It is simply not in the transformation.

import jax
import jax.numpy as jnp

def peek(x, y):
    print(type(x).__name__, x.shape, x.val.shape, x.batch_dim)
    print(type(y).__name__, y.shape)
    return x * y

jax.vmap(peek, in_axes=(1, None))(jnp.ones((3, 7, 5)), jnp.ones(5))
jax.vmap(peek, in_axes=(0, None))(jnp.ones((7, 3, 5)), jnp.ones(5))

# BatchTracer (3, 5) (3, 7, 5) 1
# ArrayImpl (5,)
# BatchTracer (3, 5) (7, 3, 5) 0
# ArrayImpl (5,)

The annotation on that integer field is NotMapped | int | RaggedAxis, and NotMapped is defined two lines above as type(None). So the None you write in in_axes and the internal marker for a value with no batch axis are the same object, aliased as not_mapped for readability. A TODO sitting above the alias asks for a real sentinel type instead.

The aval property is where the axis disappears. When batch_dim is an integer it calls core.mapped_aval, which returns the abstract value with that axis deleted, and that abstract value is what every primitive inside your function checks its shapes against. Nothing downstream has to know a batch exists.

That TODO has since been answered, in the direction of deleting the alias rather than replacing it. On jax's current main the name not_mapped does not appear in the file at all and the code compares against None directly, so if you read the modern file expecting the alias, it is gone rather than moved.

### tracer

# TODO(mattjj): use a special sentinel type rather than None
NotMapped = type(None)
not_mapped = None


class BatchTracer(Tracer):
  __slots__ = ['val', 'batch_dim', 'source_info']

  def __init__(self, trace, val, batch_dim: NotMapped | int | RaggedAxis,
               source_info: source_info_util.SourceInfo | None = None):
    if config.enable_checks.value:
      assert type(batch_dim) in (NotMapped, int, RaggedAxis)
      if type(batch_dim) is int:
        aval = core.get_aval(val)
        assert 0 <= batch_dim < len(aval.shape)
    self._trace = trace
    self.val = val
    self.batch_dim = batch_dim
    self.source_info = source_info

  @property
  def aval(self):
    aval = core.get_aval(self.val)
    if self.batch_dim is not_mapped:
      return aval
    elif type(self.batch_dim) is int:
      return core.mapped_aval(aval.shape[self.batch_dim], self.batch_dim, aval)

Every primitive your function calls lands in one function, and that function is short enough to hold in your head. It collects the values and their batch dims, then takes one of four branches, and which branch it takes is decided entirely by two dictionary lookups and one boolean.

Read the elif args_not_mapped branch before the others, because it is the one people do not expect. When no argument carries a batch dim, the primitive binds on the parent trace with the raw values, no rule is consulted, and the result is not wrapped in a tracer. Work that does not touch the batch stays outside the transformation entirely, at no cost.

The last branch is the refusal. A primitive in neither table raises NotImplementedError, with the primitive's own name in the message. Batching is opt-in per primitive, and the opt-in is a dictionary entry.

  def process_primitive(self, p, tracers, params):
    if config.dynamic_shapes.value:
      p.abstract_eval(*(map(core.get_aval, tracers)), **params)
    vals_in, dims_in = unzip2(map(self.to_batch_info, tracers))
    args_not_mapped = all(bdim is not_mapped for bdim in dims_in)
    if p in fancy_primitive_batchers:
      if (args_not_mapped
          and p in skippable_batchers
          and not any(self.axis_data.name == axis_name
                      for axis_name in skippable_batchers[p](params))):
        # no-op shortcut
        return p.bind_with_trace(self.parent_trace, vals_in, params)
      else:
        with core.set_current_trace(self.parent_trace):
          val_out, dim_out = fancy_primitive_batchers[p](self.axis_data, vals_in, dims_in, **params)
    elif args_not_mapped:
      # no-op shortcut
      return p.bind_with_trace(self.parent_trace, vals_in, params)
    elif p in primitive_batchers:
      with core.set_current_trace(self.parent_trace):
        val_out, dim_out = primitive_batchers[p](vals_in, dims_in, **params)
    else:
      raise NotImplementedError("Batching rule for '{}' not implemented".format(p))

sin does not care where the batch axis sits, because it treats every element the same. Its rule asserts that all the incoming batch dims agree, binds the primitive on the batched values unchanged, and hands the same dim back out. Three lines, and defvectorized is the one-liner that installs it.

Binary primitives need a little more, because their two operands can disagree about where the batch is or whether there is one. The broadcasting rule moves each batched operand's axis to the front and gives each unmapped operand a size-1 axis there instead, then leans on the primitive's own broadcasting to finish. A comment above that line says exactly why the inserted axis has size 1 rather than the batch size.

Reductions are the third shape of rule, and they only have to fix up a parameter. Adding a batch axis at position 0 shifts every axis in the axes parameter up by one, which is bookkeeping, not new work: reduce_sum[axes=(0,)] over a (3, 5) becomes reduce_sum[axes=(1,)] over a (4, 3, 5).

def defvectorized(prim):
  primitive_batchers[prim] = partial(vectorized_batcher, prim)

def vectorized_batcher(prim, batched_args, batch_dims, **params):
  assert all(batch_dims[0] == bd for bd in batch_dims[1:]), batch_dims
  return prim.bind(*batched_args, **params), batch_dims[0]

  else:
    # We pass size of 1 here because (1) at least one argument has a real batch
    # dimension and (2) all unmapped axes can have a singleton axis inserted and
    # then rely on the primitive's built-in broadcasting.
    args = [bdim_at_front(x, d, 1) if np.ndim(x) else x
            for x, d in zip(args, dims)]

Count the tables and there are two, holding 180 and 21 entries on this machine once one vmap has run. The counts move with what has been imported, since a rule registers when its module loads, so treat those two numbers as a reading rather than a constant.

What separates the tables is one extra argument. A rule in the second table is handed axis_data first, carrying the mapped axis's name and size, and the twenty-one primitives that need it are exactly the ones for which the identity of the axis matters: the collectives, the control-flow primitives, pjit, remat2, and the custom-derivative calls.

sin, add, dot_general, reduce_sum and transpose all sit in the first table, and so does pallas_call. A Pallas kernel is batchable because somebody wrote it a rule, not because vmap can see inside it.

Current jax has collapsed this split. On main every rule lives in the fancy table, primitive_batchers survives only as a compatibility proxy that wraps old-style rules, and defvectorized now registers into the fancy table. The two-argument shape of a rule is the one that won.

import jax
import jax.numpy as jnp
import jax._src.interpreters.batching as B

jax.vmap(lambda x: jnp.sum(jnp.tanh(x)))(jnp.ones((2, 2)))
print(len(B.primitive_batchers), len(B.fancy_primitive_batchers))
print(sorted(p.name for p in B.fancy_primitive_batchers))

# 180 21
# ['all_gather', 'all_to_all', 'axis_index', 'cond', 'custom_linear_solve',
#  'custom_vjp_call_jaxpr', 'host_local_array_to_global_array', 'pbroadcast',
#  'pgather', 'pjit', 'pmax', 'pmin', 'ppermute', 'psum', 'psum2',
#  'reduce_scatter', 'remat2', 'remat_opt', 'scan', 'sharding_constraint',
#  'while']

Define a primitive of your own with an implementation and an abstract eval and it runs, and it traces, and it lowers. Call vmap on it and the refusal arrives with the primitive's name in it, because nothing about an impl or an abstract eval implies anything about batching.

One line fixes it. defvectorized says this primitive is elementwise, which is enough for the interpreter to leave the batch axis alone, and the same call that raised a moment ago now returns a (4, 3). Look at the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → afterward and the equation is still mystery a, one primitive, operating on a batched shape.

A primitive is batchable when someone has written it a rule, and not before.

import jax
import jax.numpy as jnp
import jax._src.interpreters.batching as B
from jax._src import core as jcore

mystery_p = jcore.Primitive("mystery")
mystery_p.def_impl(lambda v: v * 2.0)
mystery_p.def_abstract_eval(lambda v: jcore.ShapedArray(v.shape, v.dtype))

print(mystery_p.bind(jnp.ones(3)))
try:
    jax.vmap(mystery_p.bind)(jnp.ones((4, 3)))
except NotImplementedError as err:
    print(err)

B.defvectorized(mystery_p)
print(jax.vmap(mystery_p.bind)(jnp.ones((4, 3))).shape)
print(jax.make_jaxpr(jax.vmap(mystery_p.bind))(jnp.ones((4, 3))))

# [2. 2. 2.]
# Batching rule for 'mystery' not implemented
# (4, 3)
# { lambda ; a:f32[4,3]. let b:f32[4,3] = mystery a in (b,) }