Skip to content

← Flax step 90 of 100

Medium Primitives

Sharding Constraints: Values and Layout

Jump to the task

Values and device placement

A sharded array still has one logical shape and one set of values. Sharding describes which devices hold those values. For example, an (8, 4) array can have one row on each of eight devices while still behaving as an (8, 4) array in your JAX code.

jax.lax.with_sharding_constraint lets you specify the layout of an intermediate array inside a compiled computation. It imposes a constraint at that point, not a suggestion or a rule for every later operation. Later operations can use different layouts, and meeting a constraint can require communication between devices. Measure the result before assuming a constraint makes your program faster.

A complete example

This example uses the local devices available to the process. It works with one device too, though then there is no work to split across devices. Each device gets two rows. The matrix multiplication maps four input features to three output features.

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

devices = np.array(jax.local_devices())
mesh = Mesh(devices, axis_names=("data",))
row_sharding = NamedSharding(mesh, P("data", None))

@jax.jit
def project(x, w):
    x = jax.lax.with_sharding_constraint(x, row_sharding)
    return x @ w

rows = 2 * len(devices)
x = jnp.arange(rows * 4, dtype=jnp.float32).reshape(rows, 4)
w = jnp.ones((4, 3), dtype=jnp.float32)
y = project(x, w)
np.testing.assert_allclose(y, np.asarray(x) @ np.asarray(w))
assert y.shape == (rows, 3)

The constraint on x preserves its shape and values. The multiplication changes the values and the final dimension. To specify a compiled function's output placement, use jax.jit(..., out_shardings=...).

NamedSharding carries its mesh explicitly. A bare PartitionSpec needs an active mesh context to give its axis names meaning. Eager calls are not generally no-ops: they can change placement too. Do not rely on omitting jit to disable a constraint or bypass its mesh requirements.

Inside jit, an intermediate can even have uneven sharding. If such an unevenly sharded value is returned, it comes out replicated. This differs from requiring a global array to have evenly sized shards at a function boundary. See the JAX API reference for the constraint's contract.

Problem

Implement with_sharding_passthrough(x): return the values of x as a one-dimensional JAX array, in row-major order. For example, [[1, 2, 3], [4, 5, 6]] becomes [1, 2, 3, 4, 5, 6].

Input: a float array x of arbitrary shape. Output: an array of shape (x.size,), preserving the element order.

The graded task is a small flattening exercise. It does not call or emulate with_sharding_constraint, and passing its tests does not verify device placement. Flattening changes the logical shape; a sharding constraint by itself does not. Use the runnable example above to explore the actual API on your machine.

Read the layout

The diagram below shows a separate two-axis example: an (8, 12) array split by P("data", "model") on a (2, 4) mesh. There are eight distinct (4, 3) blocks, with no replication. The graded function does not create this layout.

Loading visualization…