When mapped computations need values from one another—such as global normalization or shared gradient updates—use a collective operation. JAX collectives communicate across a named vmap() or pmap() axis instead of treating every mapped element or device independently. section 6.2.5 section 7.2.2
Give the mapped axis a label with axis_name, then use that same label inside jax.lax operations. For example, psum(x, 'batch') computes an all-reduce sum across the batch, letting each element divide by the same total. The pattern works with vmap() and pmap(). section 6.2.5 section 7.2.2
Common JAX collective patterns include:
- Use psum() for an all-reduce sum, pmean() for a mean, and pmin() or pmax() for global extrema. section 6.2.5
Normalize values across mapped devices
A pmap() function can normalize eight values against their global sum: norm = jax.pmap(lambda x: x / jax.lax.psum(x, axis_name='p'), axis_name='p'). Applying norm(arr) produces values whose combined sum is 1. For a larger two-dimensional array, sum within each device before calling psum() when each device owns a group of elements.
Collectives are especially useful in data-parallel training. Each device computes local gradients, psum() aggregates them, and every device applies the same update. Across hosts, each process prepares its local shard while collectives communicate across global devices; communication overhead can limit scaling. section 7.3.2 section 7.4 section 7.0
