Where the axis goes

Batch an elementwise function along axis 1 and nothing inside the function moves. sin and mul run on the (3, 4) array exactly as it arrived, because neither of them cares which axis is which, and a single transpose appears at the very end to put the batch where out_axes asked for it.

Ask for it back where it started and even that goes away. With in_axes=1 and out_axes=1 the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → is character for character the jaxpr of the unbatched function, one sin and one mul, no transpose anywhere. The rewrite added nothing because nothing needed to move.

So a move is not something vmap does to your data on the way in. It is something it does at whichever boundary the axis specs disagree about, and if they agree with each other and with the primitives, there is no boundary work to do.

>>> f = lambda x: jnp.sin(x) * 2.0
>>> jax.make_jaxpr(jax.vmap(f, in_axes=0))(jnp.ones((4, 3)))
{ lambda ; a:f32[4,3]. let b:f32[4,3] = sin a; c:f32[4,3] = mul b 2.0 in (c,) }

>>> jax.make_jaxpr(jax.vmap(f, in_axes=1))(jnp.ones((3, 4)))
{ lambda ; a:f32[3,4]. let
    b:f32[3,4] = sin a
    c:f32[3,4] = mul b 2.0
    d:f32[4,3] = transpose[permutation=(1, 0)] c
  in (d,) }

>>> jax.make_jaxpr(jax.vmap(f, in_axes=1, out_axes=1))(jnp.ones((3, 4)))
{ lambda ; a:f32[3,4]. let b:f32[3,4] = sin a; c:f32[3,4] = mul b 2.0 in (c,) }

The function that reconciles where the axis is with where you asked for it is matchaxis, and it is a chain of four cases. Same place, nothing happens. Two integers that differ, one moveaxis, which is the transpose you saw. Not mapped going to mapped, a broadcast. Mapped going to not mapped, an error, unless the caller passed sum_match, in which case the axis is summed away.

That last case is not reachable from out_axes. It exists for the transpose of a batched program, which the next lesson is about, and its presence here is the reason a broadcast in the forward direction turns into a sum in the backward one.

The error the fourth case raises is the one you meet when you write out_axes=None for something that really did get batched. JAX reformats it before you see it, and what arrives reads at vmap out_axes, got axis spec None but output was batched on axis 0.

  if src == dst:
    return x
  elif type(src) == type(dst) == int:
    return moveaxis(x, src, dst)
  elif src is not_mapped and dst is not not_mapped:
    return broadcast(x, sz, canonicalize_axis(dst, np.ndim(x) + 1))
  elif dst is not_mapped and sum_match:
    return x.sum(src)
  else:
    if (not isinstance(axis_name, core._TempAxisName) and
        axis_name is not core.no_axis_name):
      raise ValueError(f'vmap has mapped output ({axis_name=}) but out_axes is {dst}')

def matchaxis(axis_name, sz, src, dst, x, sum_match=False):
  if dst == jumble_axis:
    x = bdim_at_front(x, src, sz)
    elt_ty = x.aval.update(shape=x.shape[1:])
    aval = JumbleTy(core.Var('', core.ShapedArray((), np.dtype('int32'))),
                    x.shape[0], elt_ty)
    return Jumble(aval, x)
  try:
    _ = core.get_aval(x)
  except TypeError as e:
    raise TypeError(f"Output from batched function {x!r} with type "
                    f"{type(x)} is not a valid JAX type") from e
  if src == dst:
    return x
  elif type(src) == type(dst) == int:
    return moveaxis(x, src, dst)
  elif src is not_mapped and dst is not not_mapped:
    return broadcast(x, sz, canonicalize_axis(dst, np.ndim(x) + 1))
  elif dst is not_mapped and sum_match:
    return x.sum(src)
  else:
    if (not isinstance(axis_name, core._TempAxisName) and
        axis_name is not core.no_axis_name):
      raise ValueError(f'vmap has mapped output ({axis_name=}) but out_axes is {dst}')
    else:
      raise SpecMatchError(None, None, None)

A matmul does care which axis is which, and its rule handles that without adding an equation. Batch the left operand on axis 0 and the contraction moves from axis 0 to axis 1 of that operand, written into dimension_numbers. Batch both operands and a genuine batch-dimension pair appears in the second half of that parameter, which is what dot_general has always had for batched matmuls.

Count the equations across those variants and the count does not move. Four equations for the unbatched program, four for in_axes=0, four for in_axes=1. The rewrite is inside a parameter, not in the equation list, which is worth knowing before you go looking for the batching in a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → and fail to find it.

Reductions behave the same way one level simpler. reduce_sum[axes=(0,)] over a (3, 5) becomes reduce_sum[axes=(1,)] over a (4, 3, 5) when the batch goes in front, and axes=(0,) again when the batch goes in the middle, because the axis numbers shift around the inserted one.

>>> dot = lambda a, b: a @ b
>>> jax.make_jaxpr(dot)(jnp.ones(8), jnp.ones((8, 5)))
{ lambda ; a:f32[8] b:f32[8,5]. let
    c:f32[5] = dot_general[dimension_numbers=(([0], [0]), ([], []))] a b
  in (c,) }

>>> jax.make_jaxpr(jax.vmap(dot, in_axes=(0, None)))(jnp.ones((32, 8)), jnp.ones((8, 5)))
{ lambda ; a:f32[32,8] b:f32[8,5]. let
    c:f32[32,5] = dot_general[dimension_numbers=(([1], [0]), ([], []))] a b
  in (c,) }

>>> jax.make_jaxpr(jax.vmap(dot, in_axes=(0, 0)))(jnp.ones((32, 8)), jnp.ones((32, 8, 5)))
{ lambda ; a:f32[32,8] b:f32[32,8,5]. let
    c:f32[32,5] = dot_general[dimension_numbers=(([1], [1]), ([0], [0]))] a b
  in (c,) }

Broadcast a 1000-element weight vector across a batch of 8 and you might expect an (8, 1000) to appear in the jaxpr. What appears is a (1, 1000). The rule gives the unmapped operand a size-1 axis and lets mul do the rest, which is the comment from the previous lesson made visible: nothing is materialized per example.

Push it further and the broadcast disappears too. When an output does not depend on any batched input, the no-op shortcut keeps it unmapped end to end, so a two-output function with out_axes=(None, 0) produces a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → where the first output was computed once, at its original shape, and never widened.

The widening only happens when you ask for it. Set out_axes=0 on an unmapped output and a broadcast_in_dim shows up at the end to give it the batch axis it never had, which is matchaxis taking its third case.

>>> useboth = lambda w, x: jnp.sum(w * x)
>>> jax.make_jaxpr(jax.vmap(useboth, in_axes=(None, 0)))(jnp.ones(1000), jnp.ones((8, 1000)))
{ lambda ; a:f32[1000] b:f32[8,1000]. let
    c:f32[1,1000] = broadcast_in_dim[
      broadcast_dimensions=(np.int64(1),)
      shape=(1, 1000)
      sharding=None
    ] a
    d:f32[8,1000] = mul c b
    e:f32[8] = reduce_sum[axes=(np.int64(1),)] d
  in (e,) }

>>> k = lambda x, y: (jnp.sin(y), x * 2.0)
>>> jax.make_jaxpr(jax.vmap(k, in_axes=(0, None), out_axes=(None, 0)))(jnp.ones(4), jnp.ones(3))
{ lambda ; a:f32[4] b:f32[3]. let
    c:f32[3] = sin b
    d:f32[4] = mul a 2.0
  in (c, d) }

Eight calls of one function, mv(A, v) = A @ v, with the batch put in different places. Predict the result shape and the equation list for each before you read the right-hand columns, and check the third and fourth rows against each other: a negative in_axes counts from the right of that argument, so -2 on a rank-3 array is axis 1 and behaves identically.

The pair worth staring at is rows five and seven. Same call, same inputs, and moving out_axes from its default 0 to 1 removes an equation, because batching v on axis 0 naturally puts the batch at position 1 of the result and asking for it at position 0 is what forces the transpose. Reaching for the default out here costs a transpose that the code did not need.

Two rules on the table's edges are worth carrying. Axis integers must be in [-ndim, ndim) for the array they apply to, and keyword arguments are always mapped along axis 0 with no way to say otherwise, both of which the vmap docstring states outright. When arguments are pytrees, in_axes has to be a tree prefix of the argument tuple, so a spec may stop short and cover a whole subtree with one entry, but it may never disagree with the structure below it.

Three nested vmaps over a product of three vectors give a (2, 3, 4), the axes stacking outside in, and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → holds no nested structure of any kind. Seven equations come out, five broadcast_in_dims inserting size-1 axes and two muls. What a reader would call an outer product is what the nesting compiled to, because each layer only ever rewrote the layer below it.

Two nested maps over a matmul collapse further still, into one dot_general with dimension_numbers=(([1], [1]), ([], [])) and no batch dimensions at all. The pairwise structure lives entirely in which axes get contracted.

Swapping the nesting order does not change the result's shape here, only which axis is which. vmap(vmap(mv, in_axes=(None, 0)), in_axes=(0, None)) and the same pair reversed both return a (5, 5, 2), and one is the other with axes 0 and 1 exchanged. Shape is a weak check on a nest; if you want to know you got the order right, transpose one against the other and compare.

import jax
import jax.numpy as jnp

g = lambda a, b, c: a * b * c
three = jax.vmap(jax.vmap(jax.vmap(g, in_axes=(None, None, 0)),
                          in_axes=(None, 0, None)), in_axes=(0, None, None))
print(three(jnp.ones(2), jnp.ones(3), jnp.ones(4)).shape)

mv = lambda A, v: A @ v
A, vs = jnp.ones((5, 2, 3)), jnp.ones((5, 3))
p = jax.vmap(jax.vmap(mv, in_axes=(None, 0)), in_axes=(0, None))
q = jax.vmap(jax.vmap(mv, in_axes=(0, None)), in_axes=(None, 0))
print(p(A, vs).shape, q(A, vs).shape)
print(bool(jnp.all(p(A, vs) == jnp.swapaxes(q(A, vs), 0, 1))))

# (2, 3, 4)
# (5, 5, 2) (5, 5, 2)
# True