The output sharding is a fact you can read
Give a jitted function a sharded input and the output comes back sharded, with nothing in the body naming a mesh. Six calls on one (8, 16) input show the pattern: an elementwise op keeps both names, summing over an axis drops that axis's name, transposing swaps the two, and a matmul against a P("model", None) weight leaves the result split on data alone.
Reading the spec back is cheap enough to do as a habit. One call plus out.sharding.spec answers what the whole propagation settled on, which beats reasoning it out, because the reasoning has a hole in it: (8, 16) under P("data", "model") reshaped to (16, 8) comes back P("data"), not what the two-name input suggests.
How the compiler works any of this out belongs to 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 the propagation and the rewrite. What belongs here is the observation itself. Input shardings are the only thing you stated, and every intermediate got one anyway.
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
import jax
import jax.numpy as jnp
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")))
w = jax.device_put(jnp.ones((16, 32)), NamedSharding(mesh, P("model", None)))
for label, f, args in [("a * 2", lambda a: a * 2, (x,)),
("a.sum(0)", lambda a: a.sum(0), (x,)),
("a.sum(1)", lambda a: a.sum(1), (x,)),
("a.T", lambda a: a.T, (x,)),
("a.reshape(16, 8)", lambda a: a.reshape(16, 8), (x,)),
("a @ b", lambda a, b: a @ b, (x, w))]:
print(label.ljust(17), jax.jit(f)(*args).sharding.spec)
# a * 2 PartitionSpec('data', 'model')
# a.sum(0) PartitionSpec('model',)
# a.sum(1) PartitionSpec('data',)
# a.T PartitionSpec('model', 'data')
# a.reshape(16, 8) PartitionSpec('data',)
# a @ b PartitionSpec('data',)| call | output spec | what happened to the names |
|---|---|---|
| a * 2 | P("data", "model") | both kept: an elementwise op moves nothing |
| a.sum(0) | P("model") | the summed axis went, and its name with it |
| a.sum(1) | P("data") | the same rule on the other axis |
| a.T | P("model", "data") | the names followed their axes |
| a.reshape(16, 8) | P("data") | one name survived the new shape, the other did not |
| a @ b | P("data") | the contracted axis carried "model" and was contracted away |
In_shardings will not move a committed argument
out_shardings does what its name says. Ask for P("model", "data") on the way out and that is what comes back, whatever propagation would have chosen for you.
in_shardings behaves differently, and the difference is worth learning once rather than discovering under a deadline. On an uncommitted argument it places the array, so a plain jnp.ones gets laid across the mesh on the way in. On a committed argument whose sharding disagrees, it raises, saying the sharding passed to pjit does not match the sharding on the respective arg.
The refusal is the useful half. An input reshard is real data movement, and JAX makes you write the device_put that performs it instead of folding it invisibly into the call that was supposed to be the compiled step.
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P
mesh = jax.make_mesh((4, 2), ("data", "model"))
here = NamedSharding(mesh, P("data", "model"))
there = NamedSharding(mesh, P("model", "data"))
x = jax.device_put(jnp.ones((8, 16)), here)
print(jax.jit(lambda a: a * 2, out_shardings=there)(x).sharding.spec)
print(jax.jit(lambda a: a * 2, in_shardings=here)(jnp.ones((8, 16))).sharding.spec)
try:
jax.jit(lambda a: a * 2, in_shardings=there)(x)
except ValueError as err:
print(err)
# PartitionSpec('model', 'data')
# PartitionSpec('data', 'model')
# Sharding passed to pjit does not match the sharding on the respective arg. Got pjit
# sharding: NamedSharding(mesh=Mesh('data': 4, 'model': 2), spec=PartitionSpec('model',
# 'data'), memory_kind=unpinned_host),
# arg sharding: NamedSharding(mesh=Mesh('data': 4, 'model': 2),
# spec=PartitionSpec('data', 'model'), memory_kind=unpinned_host) for arg shape:
# float32[8,16]The compiled object tells you what it decided
Lowering and compiling ahead of the run gives you the same answers without executing anything. jax.jit(f).lower(x).compile() hands back an object whose input_shardings and output_shardings are the shardings the executable was actually built for.
This is the form of the question to put in a test. It costs a compile and no execution, it accepts ShapeDtypeStruct inputs with shardings attached so no array has to exist, and a wrong shape fails there rather than three steps into a training loop.
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
import jax
import jax.numpy as jnp
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")))
compiled = jax.jit(lambda a: a.sum(1)).lower(x).compile()
print(compiled.input_shardings[0][0].spec)
print(compiled.output_shardings.spec)
# PartitionSpec('data', 'model')
# PartitionSpec('data',)What a spec is worth, in bytes per device
The compiled object also carries a memory analysis, and that is where a spec stops being a preference and becomes a number. Two 512 by 512 float32 matrices multiplied on the eight-device mesh, fully replicated, hand each device 2097152 bytes of arguments and produce 1048576 bytes of output. Sharded P("data", None) against P(None, "model"), the same call hands each device 786432 bytes and produces 131072.
Both numbers check out by hand. Replicated, every device holds both whole matrices, which is 2 x 512 x 512 x 4 bytes. Sharded, a device holds a (128, 512) slice of the first and a (512, 256) slice of the second, so 262144 plus 524288 is 786432, and its output tile is (128, 256) at four bytes each, which is 131072.
What the analysis does not price is the traffic. Both plans compute the same product, and the sharded one buys its smaller footprint with communication. The chapter's reading list points at the cost model for that trade, and LAB·J4 is where you read which collective a given pair of specs produced.
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P
mesh = jax.make_mesh((4, 2), ("data", "model"))
matmul = jax.jit(lambda a, b: a @ b)
for spec_a, spec_b in [(P(), P()), (P("data", None), P(None, "model"))]:
A = jax.ShapeDtypeStruct((512, 512), jnp.float32, sharding=NamedSharding(mesh, spec_a))
B = jax.ShapeDtypeStruct((512, 512), jnp.float32, sharding=NamedSharding(mesh, spec_b))
m = matmul.lower(A, B).compile().memory_analysis()
print(spec_a, spec_b, m.argument_size_in_bytes, m.output_size_in_bytes)
# PartitionSpec() PartitionSpec() 2097152 1048576
# PartitionSpec('data', None) PartitionSpec(None, 'model') 786432 131072Check yourself
01 An (8, 16) input sharded P("data", "model") goes into jit(lambda a: a.sum(0)). What comes out?
A result sharded P("model"). Axis 0 was summed away, so the name that split it goes too, and the surviving axis keeps the mesh axis it already had.
02 Why does in_shardings raise on a committed argument instead of resharding it?
Because an input reshard is real data movement, and JAX will not do it silently inside the compiled call. Either device_put the array first, or hand the function an uncommitted array, which in_shardings will place.
03 Where does a per-device byte count for a given pair of specs come from?
From memory_analysis() on the object jit(f).lower(...).compile() returns. It reports argument, output and temp sizes per device, so two candidate specs can be compared before either one executes.
Readings
- jax.jit ↗ in_shardings and out_shardings, with the rules for what may be omitted
- Ahead-of-time lowering and compilation ↗ lower, compile, and the analyses hanging off the compiled object
- jax.ShapeDtypeStruct ↗ the input stand-in that carries a sharding, so nothing has to be allocated