Allreduce combines values across participating devices and returns the reduced result to every participant. In JAX, collective operations use a named mapped axis; psum, pmax, pmin, and pmean provide sum, maximum, minimum, and mean reductions. section 7.2.2
The practical pattern is to map a computation with pmap() or vmap(), assign an axis_name, and reference that name inside the function. For example, dividing each value by lax.psum(x, axis_name='p') normalizes values using the sum across the mapped axis. section 6.2.5 section 7.2.2
Normalize values across devices
arr = jnp.array(range(8)) norm = jax.pmap( lambda x: x / jax.lax.psum(x, axis_name='p'), axis_name='p' ) jnp.sum(norm(arr)) # Array(1., dtype=float32) Each mapped computation sees the same total, so the normalized outputs sum to one. For larger arrays, axis_index_groups can restrict the reduction to device groups; combining pmap() and vmap() can reduce across two named axes. section 7.2.2
Table 7.1 Collective operations (view table figure)
| Operation |
Description |
|---|---|
| all_gather(x, axis_name, *[, ...]) |
Gathers values of x across all replicas |
| all_to_all(x, axis_name, split_axis, ...[, ...]) |
Materializes the mapped axis and maps a different axis |
| psum(x, axis_name, *[, axis_index_groups]) |
Computes an all-reduce sum on x over the pmapped axis axis_name |
| psum_scatter(x, axis_name, *[, ...]) |
Like psum(x, axis_name) but each device retains only part of the result. It’s like the first part of psum without gathering the result. |
| pmax(x, axis_name, *[, axis_index_groups]) |
Computes an all-reduce max on x over the pmapped axis axis_name |
| pmin(x, axis_name, *[, axis_index_groups]) |
Computes an all-reduce min on x over the pmapped axis axis_name |
| pmean(x, axis_name, *[, axis_index_groups]) |
Computes an all-reduce mean on x over the pmapped axis axis_name |
| ppermute(x, axis_name, perm) |
Performs a collective permutation according to the permutation perm |
| pshuffle(x, axis_name, perm) |
Serves as a convenience wrapper of jax.lax.ppermute with alternate permutation encoding |
| pswapaxes(x, axis_name, axis, *[, ...]) |
Swaps the pmapped axis axis_name with the unmapped axis axis |
| axis_index(axis_name) |
Returns the index along the mapped axis axis_name |
Use the collective-operations table to compare allreduce variants with all-gather, all-to-all, and psum_scatter. section 7.2.2
