Composed with grad and jit

Put vmap outside grad and the whole program grows a batch axis, gradient equations included, and what comes back is a (4, 3) with one row per example. Put grad outside vmap and the batch axis is still there through the middle of the program, but the result is a (3,), one gradient for the batch.

The difference lands in three equations at the end of the second jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, and they are worth finding by eye: reduce_sum, then reshape, then reduce_sum again. That tail is the transpose of the broadcast that the batching rule inserted for the unmapped w on the way in, collapsing the batch axis back out of the cotangent.

Which means the two orderings are related by exactly that sum, and you can check it in one line: the rows of the per-example result add up to the batch gradient, to the last bit on this machine. If you only ever need the sum, either ordering works; if you need the rows, only one does.

>>> loss = lambda w, x: jnp.sum(jnp.tanh(x * w))
>>> jax.make_jaxpr(jax.vmap(jax.grad(loss), in_axes=(None, 0)))(jnp.ones(3), jnp.ones((4, 3)))
{ lambda ; a:f32[3] b:f32[4,3]. let
    ...
    i:f32[4,3] = mul h f
    j:f32[4,3] = mul i e
    k:f32[4,3] = add_any i j
    l:f32[4,3] = mul b k
  in (l,) }

>>> batched = lambda w, xs: jnp.sum(jax.vmap(lambda x: jnp.sum(jnp.tanh(x * w)))(xs))
>>> jax.make_jaxpr(jax.grad(batched))(jnp.ones(3), jnp.ones((4, 3)))
{ lambda ; a:f32[3] b:f32[4,3]. let
    ...
    m:f32[4,3] = mul b l
    n:f32[3] = reduce_sum[axes=(0,)] m
    o:f32[1,3] = reshape[dimensions=None new_sizes=(1, 3) sharding=None] n
    p:f32[3] = reduce_sum[axes=(np.int64(0),)] o
  in (p,) }

>>> xs = jnp.arange(12.0).reshape(4, 3)
>>> per = jax.vmap(jax.grad(loss), in_axes=(None, 0))(jnp.ones(3), xs)
>>> per.shape, per.sum(0)
((4, 3), Array([0.02974708, 0.4253622 , 0.14220996], dtype=float32))
>>> jax.grad(batched)(jnp.ones(3), xs)
Array([0.02974708, 0.4253622 , 0.14220996], dtype=float32)

That tail is not a special case anyone wrote for gradients of batched functions. It is the fourth branch of matchaxis from the last lesson, the one guarded by sum_match, doing what the transpose of a broadcast has to do.

Read it as a pair and it stops needing memorising. An unmapped input gets broadcast on the way forward, so its cotangent gets summed on the way back, because the derivative of copying a value into N places is adding the N gradients up. The batch axis you introduced with in_axes=None is exactly the axis that disappears in the reverse pass.

This also tells you where per-example gradients stop being free of assumptions. They exist because vmap sat outside grad and the sum never happened, not because JAX kept N separate tapes. There is one program, one backward pass, and an axis that was never collapsed.

pjit is one of the twenty-one primitives with a rule in the second table, and its rule batches the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → held in the equation's parameters rather than calling the equation once per example. So an inner jit does not hide its contents from the batching interpreter, and it does not force a per-example dispatch either.

The consequence is easier to state as a measurement. vmap(jit(f)) and jit(vmap(f)) trace to the same jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, text for text, on this machine. Both hold one pjit equation whose inner jaxpr is already batched, and the string comparison of the two dumps is True.

That is a useful thing to know when you inherit code with jit decorators scattered through it. Wrapping the outside in vmap does not defeat them and is not defeated by them; the batch axis goes through the call boundary and comes out the other side.

>>> f = lambda x: jnp.sin(x) * 2.0
>>> x = jnp.ones((4, 3))
>>> str(jax.make_jaxpr(jax.vmap(jax.jit(f)))(x)) == str(jax.make_jaxpr(jax.jit(jax.vmap(f)))(x))
True
>>> jax.make_jaxpr(jax.vmap(jax.jit(f)))(x)
{ lambda ; a:f32[4,3]. let
    b:f32[4,3] = pjit[
      name=<lambda>
      jaxpr={ lambda ; c:f32[4,3]. let
          d:f32[4,3] = sin c
          e:f32[4,3] = mul d 2.0
        in (e,) }
    ] a
  in (b,) }

Four errors come from the vmap machinery and nowhere else, and all four are about the axis specs rather than about your function. Every one of them fires before any tracing happens, which is why they are quick to fix: the message already contains the two things that disagree.

The first has an answer the error does not mention. Mapping nothing is legal if you say how wide the batch is, and axis_size=4 with in_axes=None runs the body once at its original shape and broadcasts the result, which is a cheap way to shape a program for a batch it does not read yet.

Two other failures look like vmap's and are not. An in_axes that names a batch of 32 against an argument whose axis is 8 raises the inconsistent-sizes error, which the museum keeps as its own exhibit. And a function whose output shape depends on its values fails under vmap for the same reason it fails under jit, since that is a rule about tracing, not about batching, and the message you get is the tracing one.

>>> jax.vmap(lambda x: x * 2.0, in_axes=None)(jnp.ones(3))
ValueError: vmap must have at least one non-None value in in_axes

>>> jax.vmap(lambda x: x * 2.0, in_axes=None, axis_size=4)(jnp.ones(3)).shape
(4, 3)
>>> jax.make_jaxpr(jax.vmap(lambda x: x * 2.0, in_axes=None, axis_size=4))(jnp.ones(3))
{ lambda ; a:f32[3]. let
    b:f32[3] = mul a 2.0
    c:f32[4,3] = broadcast_in_dim[
      broadcast_dimensions=(np.int64(1),)
      shape=(4, 3)
      sharding=None
    ] b
  in (c,) }

Everything above rests on the primitives already having rules, and the moment you add a primitive of your own that stops being true. A Pallas kernel is batchable because pallas_call is a key in the first rule table, a custom_vjp because its call primitive is a key in the second. Neither is batchable because vmap can see inside it.

So the question to ask of any new primitive is the one lesson one ended on. Does it have a rule, is that rule the elementwise one, and does the elementwise one actually describe what the primitive does to an extra axis. Get that wrong and nothing raises; you get a program that runs and computes the wrong thing.

The kernel path's Pallas arc is where that stops being hypothetical, since a kernel with a block spec has real opinions about which axis is which. The rule table is the same table.