Skip to content

← JAX Performance & Distributed step 13 of 25

Medium Primitives

NamedSharding (Replicated)

Jump to the task

Why this matters

A replicated array has a complete copy on every device in its mesh. NamedSharding(mesh, PartitionSpec()) describes that layout. Replication can make frequently used values available locally, at the cost of storing those values on each device.

This is a tradeoff, not a rule that only small constants may be replicated. Data-parallel training often replicates model parameters while splitting examples across devices. Larger models may need parameter sharding to fit. Creating or updating replicas can require communication; replicated storage does not make an entire training step communication-free.

Worked mini-example

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

mesh = Mesh(np.array(jax.devices()[:1]), ('data',))
sharding = NamedSharding(mesh, PartitionSpec())
x = jnp.array([2.0, -1.0, 4.0], dtype=jnp.float32)
replicated = jax.device_put(x, sharding)
assert replicated.sharding.is_fully_replicated
print(jnp.sum(replicated))  # 5.0

The mesh explicitly contains one device. Extra devices on the machine do not join it automatically. With a larger mesh, each device would hold all three values; the global array would still have shape (3,), and its sum would still be 5, not 5 times the device count.

Problem

Implement replicated_sharding_value(x):

  1. Create a mesh from jax.devices()[:1] with axis name 'data'.
  2. Build NamedSharding(mesh, PartitionSpec()).
  3. Place x with jax.device_put(x, sharding).
  4. Return the sum of the placed array as a float32 scalar.

x is a one-dimensional float32 JAX array. This exercise runs with one device. Its numerical tests check the sum, so they cannot distinguish this placement pattern from simply returning jnp.sum(x). Inspect the placed array's .sharding while practicing the API.

Common pitfalls

  • Replication covers the devices in the mesh, not every device in the system.
  • Placement preserves values but may copy or move data, even when the target mesh has just one device. It is not necessarily a metadata-only operation.
  • An empty PartitionSpec() replicates arrays of any rank. For a vector or matrix, PartitionSpec(None) also leaves every dimension unsplit because omitted trailing dimensions are unsplit. Use the empty spec for scalars.
  • Replicated copies represent one global array. A normal jnp.sum reduces that array, not a stack containing every physical copy.