jax.pmap takes a function and runs it in parallel across multiple devices — GPUs or TPU cores. It follows the SPMD (Single Program, Multiple Data) paradigm: every device runs the exact same program, but on different slices of the data.
import jax
import jax.numpy as jnp
# Check available devices
print(jax.devices()) # e.g., [CudaDevice(id=0), CudaDevice(id=1)]
print(jax.device_count()) # e.g., 2
# A simple function
def square(x):
return x ** 2
# pmap: run on all devices in parallel
parallel_square = jax.pmap(square)
# Input must have leading dimension = number of devices
n_devices = jax.device_count()
x = jnp.arange(n_devices * 4).reshape(n_devices, 4)
print(f"Input shape: {x.shape}") # (n_devices, 4)
# Each device gets one slice along axis 0
result = parallel_square(x)
print(f"Output shape: {result.shape}") # (n_devices, 4)
print(result)
The key concept: the leading axis of every input is the device axis. If you have 4 devices, inputs must have their first dimension equal to 4. Each device receives one slice.
import jax
import jax.numpy as jnp
# Preparing data for pmap: reshape to (n_devices, per_device_batch, ...)
n_devices = jax.device_count()
total_batch = 128 # Total examples
per_device = total_batch // n_devices
# Original data: (128, 784)
data = jax.random.normal(jax.random.PRNGKey(0), (total_batch, 784))
# Reshape for pmap: (n_devices, per_device_batch, 784)
sharded_data = data.reshape(n_devices, per_device, 784)
print(f"Sharded shape: {sharded_data.shape}")
💡 Why This Matters
pmap is the simplest way to use multiple GPUs or TPU cores in JAX. It handles data distribution, parallel execution, and result collection. For data-parallel training (the most common multi-device pattern), pmap lets you scale from 1 to N devices with minimal code changes.
Note on rank behavior: In JAX 0.9+, pmap uses pmap_no_rank_reduction behavior — the function mapped by pmap sees arrays with the same rank as the input (the device axis is preserved, not squeezed). This means if your input has shape (n_devices, batch, features), the function inside pmap sees (1, batch, features) rather than (batch, features). Keep this in mind when migrating older code.