How can a collective operation know which mapped dimension to use? In JAX, axis_name is a string label attached to an axis created by vmap() or pmap(). A collective operation uses the same label to communicate across batch elements or devices, rather than relying only on a positional dimension. section 6.2.5 section 7.2.2
Use the label to:
- Associate a mapped axis with a readable name such as batch or p. section 6.2.5
- Refer to that axis from collectives such as psum(), pmean(), pmin(), pmax(), and all_gather(). section 6.2.5 section 7.2.2
- Distinguish multiple mapped axes, including nested vmap() and pmap() transformations. section 7.2.2
- Restrict communication to groups with axis_index_groups when a full-axis reduction is not wanted. section 7.2.2
Normalize values across a vmap() axis
arr = jnp.array(range(50)) norm = jax.vmap( lambda x: x / jax.lax.psum(x, axis_name='batch'), axis_name='batch' ) result = norm(arr) Here, vmap() names its mapped axis batch, and psum() uses that same name to compute the sum across the batch. Each element divides by the shared sum, so the resulting values sum to 1.0. The identical pattern works with pmap(), where the label identifies the device axis instead.
