Modern hardware often includes multiple accelerators like GPUs or TPUs. While jax.jit optimizes code for a single device and jax.vmap handles batching efficiently, they don't inherently distribute computation across multiple devices. This chapter introduces jax.pmap (parallel map), JAX's function transformation for distributing computations by running them in parallel on different devices.

You will learn:

• The concept of Single Program, Multiple Data (SPMD) parallelism, the execution model pmap uses.
• How to use jax.pmap to execute the same function simultaneously on multiple devices, each operating on a different slice of data.
• Techniques for specifying how data should be distributed to and gathered from devices using arguments like in_axes and out_axes.
• The use of collective operations (such as sum or mean across devices using jax.lax primitives) within pmap-transformed functions.
• How pmap interacts with other JAX transformations like jit and grad.

By the end of this chapter, you will be able to apply pmap to implement data parallelism for your JAX programs on multi-device systems.