Pinning a layout, and where this course stops

Put with_sharding_constraint in the middle of a function and the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → grows an equation for it. sharding_constraint is a primitive the same way mul and reduce_sum are, and it carries the whole NamedSharding in its parameters, plus a set of unconstrained dimensions that section three fills in.

What it constrains is a value, not a variable and not a function. It says the value bound at this point is laid out this way, and everything downstream propagates from there.

Which is why the second run below matters more than it looks. Pinning a * 2 to P("data", None) changes how that intermediate is stored, and the sum after it still comes back P("data"), exactly as it does without the constraint. A constraint controls the value you named. It does not dictate the result.

{ lambda ; a:f32[8,16]. let
    b:f32[8,16] = mul a 2.0
    c:f32[8,16] = sharding_constraint[
      layout=None
      resource_env=ResourceEnv(mesh=Mesh())
      sharding=NamedSharding(mesh=Mesh('data': 4, 'model': 2), spec=PartitionSpec('data', None), memory_kind=unpinned_host)
      unconstrained_dims=set()
    ] b
    d:f32[8] = reduce_sum[axes=(1,)] c
  in (d,) }

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.lax import with_sharding_constraint
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.ones((8, 16)), NamedSharding(mesh, P("data", "model")))
rows = NamedSharding(mesh, P("data", None))

def f(a):
    b = a * 2
    b = with_sharding_constraint(b, rows)
    return b.sum(1)

print(jax.make_jaxpr(f)(x))

plain = jax.jit(lambda a: (a * 2).sum(1))
pinned = jax.jit(f)
print(plain(x).sharding.spec, pinned(x).sharding.spec)

# { lambda ; a:f32[8,16]. let
#     b:f32[8,16] = mul a 2.0
#     c:f32[8,16] = sharding_constraint[
#       layout=None
#       resource_env=ResourceEnv(mesh=Mesh())
#       sharding=NamedSharding(mesh=Mesh('data': 4, 'model': 2), spec=PartitionSpec('data', None), memory_kind=unpinned_host)
#       unconstrained_dims=set()
#     ] b
#     d:f32[8] = reduce_sum[axes=(1,)] c
#   in (d,) }
# PartitionSpec('data',) PartitionSpec('data',)

Lower the same function and the constraint is still there, as stablehlo.custom_call @Sharding sitting between the multiply and the reduce, with an mhlo.sharding string attached. The function's argument carries a string of its own in the signature, in the same notation.

Read that notation once and it stops being noise. {devices=[4,2]<=[8]} says the array is tiled four by two over the eight devices taken in order. <=[4,2]T(1,0) says take that same four by two grid transposed, which is what naming the axes the other way round produces. And last_tile_dim_replicate marks a trailing tile dimension that is not an array axis at all, so {devices=[4,1,2]<=[8] last_tile_dim_replicate} reads as four tiles down axis 0, one across axis 1, each tile held by two devices.

The table below is the whole vocabulary for a two-axis array on this mesh, measured one spec at a time. The printer that produces these strings lives in XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s hlo_sharding.cc, which settles any case a table this size does not cover.

module @jit_f attributes {mhlo.num_partitions = 8 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<8x16xf32> {mhlo.sharding = "{devices=[4,2]<=[8]}"}) -> (tensor<8xf32> {jax.result_info = ""}) {
    %cst = stablehlo.constant dense<2.000000e+00> : tensor<f32>
    %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<8x16xf32>
    %1 = stablehlo.multiply %arg0, %0 : tensor<8x16xf32>
    %2 = stablehlo.custom_call @Sharding(%1) {backend_config = "", mhlo.sharding = "{devices=[4,1,2]<=[8] last_tile_dim_replicate}"} : (tensor<8x16xf32>) -> tensor<8x16xf32>
    %cst_0 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %3 = stablehlo.reduce(%2 init: %cst_0) applies stablehlo.add across dimensions = [1] : (tensor<8x16xf32>, tensor<f32>) -> tensor<8xf32>
    return %3 : tensor<8xf32>
  }
}

P.UNCONSTRAINED in an entry says you have an opinion about the other axes and none about this one. The annotation records it as backend_config = "unspecified_dims=[1]", and propagation fills the gap, so constraining axis 0 to data and leaving axis 1 open returns a result still split P("data", "model").

Set that against pinning both axes to nothing. P() in the middle of a function forces the value to be replicated at that point, the annotation becomes {replicated}, and the output comes back P(). One of the two leaves a decision to the compiler; the other takes it away.

The distinction is easy to lose because None and UNCONSTRAINED sit in the same slot. None is a request: keep this axis whole on every device. UNCONSTRAINED is the absence of a request, and an axis you left alone can come back split.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.lax import with_sharding_constraint
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.ones((8, 16)), NamedSharding(mesh, P("data", "model")))

def pinned(a):
    return with_sharding_constraint(a * 2, NamedSharding(mesh, P("data", P.UNCONSTRAINED)))

def gathered(a):
    return with_sharding_constraint(a * 2, NamedSharding(mesh, P()))

for f in (pinned, gathered):
    line = [l for l in jax.jit(f).lower(x).as_text().splitlines() if "@Sharding" in l][0]
    print(line.strip())
    print("   out:", jax.jit(f)(x).sharding.spec)

# %2 = stablehlo.custom_call @Sharding(%1) {backend_config = "unspecified_dims=[1]",
#   mhlo.sharding = "{devices=[4,1,2]<=[8] last_tile_dim_replicate}"} :
#   (tensor<8x16xf32>) -> tensor<8x16xf32>
#    out: PartitionSpec('data', 'model')
# %2 = stablehlo.custom_call @Sharding(%1) {backend_config = "",
#   mhlo.sharding = "{replicated}"} : (tensor<8x16xf32>) -> tensor<8x16xf32>
#    out: PartitionSpec()

Ask for something impossible and nothing raises. Slicing six rows out of the array and constraining that to P("data", None) asks for six rows split four ways, which has no answer, and the call still returns.

The annotation goes into the module exactly as written, {devices=[4,1,2]<=[8] last_tile_dim_replicate} on a tensor<6x16xf32>. The executable that comes back holds the result replicated, P(), with a (6, 16) shard on every device. The constraint was carried to the compiler, and the compiler declined it.

So a constraint is not an assertion, and treating it as one is how a program ends up quietly replicating a value it meant to split. The check that catches this is the one from lesson three: read the output sharding back, or read the compiled executable's, and compare it with what you asked for.

A constraint states an intent. Only the compiled program states an outcome.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.lax import with_sharding_constraint
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.ones((8, 16)), NamedSharding(mesh, P("data", "model")))

# six rows cannot be split four ways
f = jax.jit(lambda a: with_sharding_constraint(a[:6], NamedSharding(mesh, P("data", None))))
line = [l for l in f.lower(x).as_text().splitlines() if "@Sharding" in l][0]
print(line.strip())
print(f(x).sharding.spec, f(x).addressable_shards[0].data.shape)

# %1 = stablehlo.custom_call @Sharding(%0) {backend_config = "",
#   mhlo.sharding = "{devices=[4,1,2]<=[8] last_tile_dim_replicate}"} :
#   (tensor<6x16xf32>) -> tensor<6x16xf32>
# PartitionSpec() (6, 16)

Calling with_sharding_constraint on a real array, outside any trace, does not raise on this version. The array comes back with the spec you asked for, committed, with the shard shape that spec implies, which is device_put in different clothing.

That convenience hides a real difference. Inside a trace there is a value the compiler has not laid out yet, and the call annotates it. Outside, there is nothing left to annotate, so the call moves data, with everything lesson two said about what a reshard costs.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.lax import with_sharding_constraint
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.ones((8, 16)), NamedSharding(mesh, P("data", "model")))

y = with_sharding_constraint(x, NamedSharding(mesh, P(None, "model")))
print(y.sharding.spec, y.committed, y.addressable_shards[0].data.shape)

# PartitionSpec(None, 'model') True (8, 8)

Four questions about sharding have separate homes, and two of them are here. What a mesh and a spec mean, and what device_put does with them, is lesson one and lesson two. What jit propagates and what a spec costs per device is lesson three. Chapter 10 adds the level below both, shard_map, where the function runs once per device and you name the collectives yourself.

How the propagation and the rewrite actually happen is the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → path's SPMD chapter, which reads the pass that does it. What a collective is as an agreement between devices, and how that agreement deadlocks, is its collectives chapter. The kernel path's distributed stage builds one out of raw remote DMAsA chip pushes a buffer straight into a neighbor’s memory and signals a semaphore, while its compute keeps working. The native distributed operation.taught in /l/ici →, which is the same operation three layers down.

Which spec to choose is a cost question, and the scaling book on the chapter's reading list carries that arithmetic. LAB·J4 is the runnable boundary between the two: eight devices on one laptop, and the question of which collective a given pair of specs produced, which is the first question this vocabulary makes it possible to ask.