What does jax.shard_map actually do to your function?

It calls your function once, on abstract values whose shapes are already divided down to one device. The docstring puts it as mapping a function over shards of data using a mesh of devices, where each application of f takes as input a shard of the mapped-over arguments and produces a shard of the output. Concretely: every argument aval is rewritten by dividing each dimension by the product of the mesh axis sizes its in_specs entry names, the mesh is re-typed so those axes are Manual, your Python runs once against that smaller shape, and every output aval is multiplied back up by whatever its out_specs entry names. The body becomes the jaxpr parameter of a single shard_map equation, and the whole equation lowers to one region the partitioner is told not to touch.

What is the difference between in_specs and out_specs in shard_map?

They are promises pointing in opposite directions. An in_specs entry is a division: name a mesh axis at position i and the i-th dimension of that argument is split across that axis, so the body sees a smaller array. An out_specs entry is a concatenation: name a mesh axis at position i and the shards produced along that axis are joined back along dimension i. The asymmetry is in what silence means. Leaving an axis out of in_specs means every device along it gets the same data. Leaving it out of out_specs is a claim by you that the outputs are already equal along that axis, so one value is taken rather than many concatenated, and JAX checks that claim before it will build the program.

What is check_vma in shard_map and when do you turn it off?

check_vma defaults to True and switches on a type-level record, carried on every abstract value inside the body, of which manual mesh axes that value is allowed to differ along. An argument starts out varying over exactly the axes its in_specs entry names. Ordinary primitives require their inputs to agree and pass the set through. A psum over an axis removes it. At the end, any axis your out_specs entry does not mention must not be in that set, or shard_map raises rather than emitting a program. Turn it off when the check refuses something you know is correct, which usually means the varying set is over-approximated somewhere; the docstring calls it additional validity checks and automatic differentiation optimizations, so you also give up the second thing when you disable the first.

Which collectives can you call inside a shard_map body?

The named-axis collectives in jax.lax: psum, pmean, pmax, pmin, all_gather, all_to_all, psum_scatter, ppermute, pswapaxes, pshuffle, ragged_all_to_all, and axis_index. Each one first checks that the axis name you passed actually exists in the axis environment, which is the check that produces the unbound axis name error when you call one outside a shard_map. Then, with check_vma on, the collective checks the value it is given against the varying record: a variant-to-invariant collective like psum requires the axis to be present and removes it, while the rule shared by the rest requires every axis it acts over to already be there. axis_index goes the other way, producing a value that varies over the axis you name.

What does shard_map lower to?

One region, in whichever dialect is switched on. With the Shardy partitioner, which is the default at this commit, the rule builds a single sdy.ManualComputationOp whose operands carry the in_specs shardings, whose results carry the out_specs shardings, and whose ManualAxesAttr lists exactly the mesh axes the region is manual over; the body jaxpr is lowered into the block inside it. On the older GSPMD path each operand is wrapped in a Sharding custom call and then a SPMDFullToShardShape custom call, the body is lowered as an ordinary called computation, and each result is wrapped back through SPMDShardToFullShape. There is one shortcut: if every newly manual axis has size one, no region is emitted at all and the body is lowered inline.

Why does shard_map say it could not infer replication over any axes?

Because your out_specs left a mesh axis unmentioned and the body could not prove the output is the same along it. Unmentioned means replicated, so shard_map compares the axes your spec omits against the varying record it computed for that output. If an omitted axis is still in the record, it raises a ValueError naming the spec, the mesh shape, the axis, and the suggestion to either revise the out_specs entry or pass check_vma=False. The usual fix is neither: it is a missing collective. If you meant the result to be identical on every device along that axis, something in the body has to make it identical, which is a psum, a pmean, or an all_gather over that axis.

shard_map Internals: The Per-Device Body and the Checks Around It

Everywhere else in JAX you write the whole-array program and let the partitioner work out the communication. shard_map inverts that: you write the program one device runs, on one shard, and you write the collectives yourself. This page reads the file that makes that work at one pinned commit of jax-ml/jax. What in_specs and out_specs each promise, the six checks that run before your function is ever called, the shape arithmetic that turns a global array into a per-device one and back, the type-level record of which mesh axes a value is allowed to differ along, which collectives are legal inside the body and what each one does to that record, and the single region the whole thing lowers to. The spec checker is on the page: pick a mesh, an in_specs, a body and an out_specs, and read the per-device shapes and the real verdict.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of shard_map Internals: The Per-Device Body and the Checks Around It — built to be understood, not skimmed.

Questions

What does jax.shard_map actually do to your function?: It calls your function once, on abstract values whose shapes are already divided down to one device. The docstring puts it as mapping a function over shards of data using a mesh of devices, where each application of f takes as input a shard of the mapped-over arguments and produces a shard of the output. Concretely: every argument aval is rewritten by dividing each dimension by the product of the mesh axis sizes its in_specs entry names, the mesh is re-typed so those axes are Manual, your Python runs once against that smaller shape, and every output aval is multiplied back up by whatever its out_specs entry names. The body becomes the jaxpr parameter of a single shard_map equation, and the whole equation lowers to one region the partitioner is told not to touch.
What is the difference between in_specs and out_specs in shard_map?: They are promises pointing in opposite directions. An in_specs entry is a division: name a mesh axis at position i and the i-th dimension of that argument is split across that axis, so the body sees a smaller array. An out_specs entry is a concatenation: name a mesh axis at position i and the shards produced along that axis are joined back along dimension i. The asymmetry is in what silence means. Leaving an axis out of in_specs means every device along it gets the same data. Leaving it out of out_specs is a claim by you that the outputs are already equal along that axis, so one value is taken rather than many concatenated, and JAX checks that claim before it will build the program.
What is check_vma in shard_map and when do you turn it off?: check_vma defaults to True and switches on a type-level record, carried on every abstract value inside the body, of which manual mesh axes that value is allowed to differ along. An argument starts out varying over exactly the axes its in_specs entry names. Ordinary primitives require their inputs to agree and pass the set through. A psum over an axis removes it. At the end, any axis your out_specs entry does not mention must not be in that set, or shard_map raises rather than emitting a program. Turn it off when the check refuses something you know is correct, which usually means the varying set is over-approximated somewhere; the docstring calls it additional validity checks and automatic differentiation optimizations, so you also give up the second thing when you disable the first.
Which collectives can you call inside a shard_map body?: The named-axis collectives in jax.lax: psum, pmean, pmax, pmin, all_gather, all_to_all, psum_scatter, ppermute, pswapaxes, pshuffle, ragged_all_to_all, and axis_index. Each one first checks that the axis name you passed actually exists in the axis environment, which is the check that produces the unbound axis name error when you call one outside a shard_map. Then, with check_vma on, the collective checks the value it is given against the varying record: a variant-to-invariant collective like psum requires the axis to be present and removes it, while the rule shared by the rest requires every axis it acts over to already be there. axis_index goes the other way, producing a value that varies over the axis you name.
What does shard_map lower to?: One region, in whichever dialect is switched on. With the Shardy partitioner, which is the default at this commit, the rule builds a single sdy.ManualComputationOp whose operands carry the in_specs shardings, whose results carry the out_specs shardings, and whose ManualAxesAttr lists exactly the mesh axes the region is manual over; the body jaxpr is lowered into the block inside it. On the older GSPMD path each operand is wrapped in a Sharding custom call and then a SPMDFullToShardShape custom call, the body is lowered as an ordinary called computation, and each result is wrapped back through SPMDShardToFullShape. There is one shortcut: if every newly manual axis has size one, no region is emitted at all and the body is lowered inline.
Why does shard_map say it could not infer replication over any axes?: Because your out_specs left a mesh axis unmentioned and the body could not prove the output is the same along it. Unmentioned means replicated, so shard_map compares the axes your spec omits against the varying record it computed for that output. If an omitted axis is still in the record, it raises a ValueError naming the spec, the mesh shape, the axis, and the suggestion to either revise the out_specs entry or pass check_vma=False. The usual fix is neither: it is a missing collective. If you meant the result to be identical on every device along that axis, something in the body has to make it identical, which is a psum, a pmean, or an all_gather over that axis.

shard_map Internals: The Per-Device Body and the Checks Around It

Questions

Related explainers