While pmap works well for simple data parallelism, modern JAX (0.5+) provides more flexible sharding APIs that handle complex parallelism strategies — including model parallelism, pipeline parallelism, and mixed strategies.
The Modern Sharding API
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
# Since JAX 0.7+, you can also use: jax.P as a shortcut for PartitionSpec
# Step 1: Create a mesh — a logical arrangement of devices
devices = jax.devices() # All available devices
mesh = jax.make_mesh((len(devices),), ('data',))
# Step 2: Define sharding with PartitionSpec
# P('data') means "shard along the 'data' axis of the mesh"
data_sharding = NamedSharding(mesh, P('data'))
replicated = NamedSharding(mesh, P()) # No sharding = replicated
# Step 3: Place data on devices
X = jax.random.normal(jax.random.PRNGKey(0), (128, 4))
X_sharded = jax.device_put(X, data_sharding)
# Check the sharding
jax.debug.visualize_array_sharding(X_sharded)
2D Meshes for Combined Data + Model Parallelism
import jax
import jax.numpy as jnp
# Create a 2D mesh: (data_parallel, model_parallel)
# For example, 8 devices arranged as 4x2
mesh = jax.make_mesh((4, 2), ('dp', 'mp'))
# Data is sharded along the 'dp' axis
data_sharding = jax.NamedSharding(mesh, jax.P('dp', None))
# Model weights are sharded along the 'mp' axis
weight_sharding = jax.NamedSharding(mesh, jax.P(None, 'mp'))
# With jit + sharding constraints, JAX handles communication automatically
@jax.jit
def forward(params, x):
return x @ params
# JAX inserts the necessary all-gathers and reduce-scatters
jax.P: The Convenient Alias (JAX 0.7+)
import jax
import jax.numpy as jnp
# Since JAX 0.7, jax.P is an alias for jax.sharding.PartitionSpec
mesh = jax.make_mesh((4, 2), ('dp', 'mp'))
# These are equivalent:
sharding1 = jax.NamedSharding(mesh, jax.P('dp', None))
sharding2 = jax.NamedSharding(mesh, jax.sharding.PartitionSpec('dp', None))
# jax.P is more concise and is the recommended style
When to Use pmap vs Modern Sharding
| Use pmap when... | Use jax.sharding when... |
|---|---|
| Simple data parallelism | Combined data + model parallelism |
| Small number of devices (2-8) | Large device meshes (TPU pods) |
| You want explicit control | You want the compiler to optimize communication |
| Prototyping multi-device code | Production distributed training |
💡 Why This Matters
The modern sharding API (Mesh, PartitionSpec, NamedSharding) is where JAX's multi-device story is heading. It's more composable, works with jit instead of requiring a separate pmap, and lets XLA's Shardy compiler automatically insert the optimal communication patterns. For new projects, especially those targeting TPU pods or multi-node GPU clusters, the sharding API is the recommended approach.