Skip to content

← JAX Performance & Distributed step 14 of 25

Hard Primitives

shard_map Basics

Jump to the task

Why this matters

jax.shard_map lets you write a function for each device's local block of an array. The mesh identifies devices, in_specs describes their input blocks, and out_specs describes how local results form a global array.

Unlike a mapped axis in vmap, a partitioned dimension stays in the local array's shape. Splitting a vector of length eight across four devices gives each function a vector of length two, not a scalar.

Worked mini-example

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:1]), ('data',))

double = jax.shard_map(
    lambda block: 2 * block,
    mesh=mesh, in_specs=P('data'), out_specs=P('data'),
)
x = jnp.array([2.0, -1.0, 0.5], dtype=jnp.float32)
print(double(x))  # [ 4. -2.  1.]

This mesh contains one device, so its local block is the whole vector. On a larger mesh, P('data') assembles the local outputs along axis zero. The result can remain distributed; returning it does not gather the whole array onto the host or onto every device.

Problem

Implement shard_map_double(x). Create a mesh from jax.devices()[:1] with axis name 'data'. Use jax.shard_map to double each local block, with PartitionSpec('data') for both input and output. Return its result on x.

x is a one-dimensional float32 JAX array. Return the same shape and dtype, with every value doubled. The numerical tests check values, not use of shard_map. Simply multiplying x would produce the same numerical answer.

Common pitfalls

  • Always select just the first device for this exercise. Using all devices requires a vector length divisible by the size of the data mesh axis.
  • The body sees a local block. Access to other blocks requires communication, such as a collective operation, rather than global indexing.
  • A spec may omit trailing unsplit array dimensions. It must not name a partitioned dimension that the input or output does not have.
  • out_specs describes assembly, not a reduction. Use a collective when separate local values need to be combined arithmetically.
  • Use the current jax.shard_map API. Older material imports it from jax.experimental.shard_map.

The diagram below is a separate matrix split across two mesh axes. The function above doubles a vector on a one-device mesh.

Loading visualization…