Skip to content

← JAX Performance & Distributed step 16 of 25

Hard Primitives

shard_map vs vmap

Jump to the task

Why this matters

vmap batches a function over an array axis. shard_map assigns blocks to devices and runs a function on each local block. They can be used together: shard_map distributes rows, then vmap applies a row function within each block. Neither transformation alone promises a particular speedup.

A vmap function does not need a mesh. That does not limit it to one device: it can participate in a compiled computation with sharded arrays. The choice here is between automatic batching and explicitly written local block computations, not between single-device and multi-device capability.

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',))

def row_score(row):
    return jnp.sum(row ** 2)

local_scores = jax.vmap(row_score)
mapped_scores = jax.shard_map(
    local_scores, mesh=mesh,
    in_specs=P('data'), out_specs=P('data'),
)
x = jnp.array([[2.0, -1.0], [0.0, 3.0]], dtype=jnp.float32)
print(local_scores(x))   # [5. 9.]
print(mapped_scores(x))  # [5. 9.]

Both methods compute the same per-row quantity. They also agree on a larger compatible mesh, apart from possible floating-point differences. This example explicitly selects one device, so it does not measure a distributed speedup.

Problem

Implement shard_map_vs_vmap(x) for a float32 matrix of shape (B, d):

  1. Build a mesh from jax.devices()[:1], with axis name 'data'.
  2. In the shard_map path, apply jax.vmap(lambda r: jnp.sum(r ** 2)) to each local block of rows. Use PartitionSpec('data') for input and output.
  3. In the direct vmap path, apply that row function to x.
  4. Return jnp.stack([sm_result, vmap_result], axis=0).

The float32 output has shape (2, B). Each row contains the same B scores. Numerical fixtures cannot prove both transformations were used; duplicating one path's values could pass those checks.

Common pitfalls

  • For an input (B, d) split across k data devices, a local block has shape (B / k, d). The body receives rows, not a single row.
  • PartitionSpec('data') leaves the feature dimension unsplit. This does not mean every device receives the entire matrix.
  • On a larger mesh, B must be divisible by the assigned data-axis size.
  • Stack the two length-B outputs in the stated order. Concatenating them would produce a vector of length 2B instead of a (2, B) matrix.

The diagram below is a separate two-axis mesh: row blocks repeat across the unused model axis. This exercise uses a one-device data mesh.

Loading visualization…