8 / 8 JAX & XLA Mental Model

Manual Mesh Control

You have mesh = Mesh(devices, ('data', 'model')) and want each device to run identical per-shard code with explicit collective ops (e.g. an explicit psum) rather than letting the compiler infer inter-device communication from a global-shape program. Which JAX API gives you this style?

A

jax.pmap with axis_name bound to one of the mesh's axis names

B

XLA's auto-spmd-partitioning compiler flag applied to a plain jax.jit call

C

jax.experimental.shard_map.shard_map, which maps the function over the mesh so each invocation sees only its local shard and issues explicit collectives

D

jax.jit with only in_shardings=NamedSharding(mesh, PartitionSpec('data','model')) set, since that already forces explicit per-device code

E

jax.jit(donate_argnums=(0,)) combined with a NamedSharding annotation on the donated argument