Skip to content

← JAX Performance & Distributed step 20 of 25

Hard Primitives

with_sharding_constraint

Jump to the task

Why this matters

A compiled computation can use different layouts for its inputs and its intermediate arrays. jax.lax.with_sharding_constraint(x, sharding) lets you specify the layout at a particular point inside that computation. For the GSPMD partitioner, this is a strict constraint, not a suggestion.

For example, an operation may need a complete copy of an intermediate on each device after earlier operations used partitioned data. Enforcing that layout may require communication. A constraint does not guarantee a speedup.

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',))
replicated = NamedSharding(mesh, PartitionSpec())

@jax.jit
def sum_with_constraint(x):
    constrained = jax.lax.with_sharding_constraint(x, replicated)
    return jnp.sum(constrained)

x = jnp.array([2.0, -1.0, 4.0], dtype=jnp.float32)
print(sum_with_constraint(x))  # 5.0

This example uses one device so it runs without distributed hardware. It shows where the constraint belongs, but cannot demonstrate communication between devices. The constraint preserves the array's mathematical values.

Problem

Implement constrained_sharding(x):

  1. Build a mesh from jax.devices()[:1], with axis name 'data'.
  2. Create NamedSharding(mesh, PartitionSpec()) for full replication.
  3. Apply jax.lax.with_sharding_constraint(x, sharding).
  4. Return jnp.sum of the constrained array.

x is a one-dimensional float32 JAX array. Return its sum as a float32 scalar. The function should work when called directly and when wrapped in jax.jit; you do not need to add a decorator to the submitted function.

The numerical tests check the sum, not whether the constraint was used. To study compilation, trace your function with jax.make_jaxpr and inspect its sharding constraint separately from its numerical result.

Common pitfalls

  • Keep the returned constrained value and use it in the next operation. Calling the function and then continuing with the original x misses the intended connection in the computation.
  • Do not assume the call is a no-op outside an explicit jax.jit wrapper. Eager execution can change the array's sharding and move data.
  • PartitionSpec() replicates across the given mesh. This exercise's mesh contains only the first device, even on a machine with more devices.
  • An intermediate constraint does not specify the layout of every later result. Use a compiled function's out_shardings when you need to specify its output layout.