JAX 0.10.0 Kept pmap — but Replaced the Implementation Underneath It
How the same API came to rest on a different contract for sharding and distributed execution
The JAX 0.10.0 release notes contain an unusual combination of changes.
jax.pmap still exists.
But the old C++ pmap infrastructure has been removed.
The API name remained while the dedicated parallel execution path behind it disappeared.
Today, jax.pmap is implemented in terms of jax.jit and jax.shard_map.
That means code can still look like an ordinary pmap call while relying on a different set of rules for device placement, sharding, and distributed execution.
Why would JAX remove a parallelization path that many users already relied on?
And where can the difference become visible?
What pmap has traditionally done
Large machine-learning workloads often require more than one GPU or TPU.
Data has to be divided across devices, and the same function has to run on each piece in parallel.
jax.pmap made this relatively easy to express.
A leading array axis could be mapped across devices, with each device computing on its own slice.
Collective operations such as psum could then combine information across those devices.
From the user’s perspective, many details of device topology and data placement were hidden behind a compact transformation.
Behind that convenience was a dedicated implementation built around components such as PmapSharding, PmapExecutable, and PmapFunction.
The name stayed, but the implementation changed
Starting with JAX 0.8.0, the default pmap implementation moved to a composition of jit and shard_map.
For a limited period, users who encountered regressions could switch back to the old implementation with a configuration flag.
That escape hatch was temporary.
The flag was deprecated in JAX 0.9.0 and removed in JAX 0.10.0.
At the same time, the release removed several pieces of the old pmap stack, including:
jax.sharding.PmapShardingthe C++
PmapFunctionandpmapAPIsPmapExecutableand relatedpxlainterfaceslegacy types such as
ShardingSpec,Chunked, andReplicated
pmap remains available, but it now runs on JAX’s general sharding machinery rather than a separate C++ parallel execution system.
Why remove a dedicated parallel engine?
The old design gave pmap its own execution model.
Elsewhere, JAX was moving toward a more unified model based on jit, Mesh, NamedSharding, and explicit device mappings.
Both systems handled distributed arrays, but they did not always express placement and execution in the same way.
The new implementation brings those worlds together.
shard_map describes per-device computation.
jit compiles it.
A Mesh describes the logical device topology, while PartitionSpec and NamedSharding describe how array dimensions map onto that topology.
JAX’s migration guide gives two main reasons for the change: better integration with JAX sharding and a simpler implementation.
That direction is internally coherent.
But simplifying the implementation does not preserve every implicit assumption made by existing code.
Silent resharding may no longer happen
Under the old pmap, inputs whose placement did not exactly match the computation could sometimes be adjusted behind the scenes.
The new implementation is stricter.
If an input is placed on incompatible devices or carries a sharding that does not match what the computation expects, JAX may raise an error instead of silently resharding it.
The long-term fix is to make the placement explicit.
Users may need to call jax.device_put with an appropriate NamedSharding, or change the function’s device and axis configuration to match the actual operands.
This can feel less convenient.
But implicit data movement can hide communication and performance costs.
The new boundary makes those costs visible.
Placement that JAX previously corrected for you
may now have to be stated directly in the program.
Nested pmap no longer means the same thing
Some existing programs used nested pmap calls to express more than one axis of parallelism.
That pattern is not supported in the same way by the new implementation.
An outer pmap does not have enough information about the mesh axes created by an inner pmap.
JAX recommends expressing multiple parallel axes in a single shard_map, or using explicit mesh-based parallelism instead.
The responsibility has moved.
What used to be expressed by wrapping one transformation inside another is now represented as an explicit device topology.
This is not merely a syntax change.
It changes where the program declares the structure of parallel execution.
A simple x[0] can trigger communication
One of the more revealing differences appears in a very ordinary operation.
Distributed training code often replicates a value across devices and later reads the first copy for logging or checkpointing:
x[0]
A user may assume that every device holds the same value, so reading the first one should be cheap.
The old PmapSharding implementation had a special fast path that could return data from a single device.
NamedSharding follows the logical semantics of the global array more strictly.
If the leading dimension is physically sharded and the result of x[0] is expected to be available with replicated semantics, JAX may need to communicate across devices.
In a particular leading-axis sharding pattern, it can gather the distributed array and then take the requested slice.
The source code still contains a simple indexing operation.
But the execution meaning underneath it has changed.
When only a local shard is needed for logging or checkpointing, accessing addressable_shards[0].data can avoid that global operation.
Multi-host programs expose the boundary more clearly
In a multi-host JAX program, each process directly owns only its local data.
The new pmap implementation converts host-local arrays into the global array representation expected by shard_map, then converts the outputs back to host-local form.
The official migration guide notes that this round trip cannot always be avoided.
If its cost becomes significant, JAX recommends migrating directly to shard_map.
A program that appears unchanged on one host may therefore expose different communication and synchronization costs when scaled across machines.
The new model treats the cluster as a more unified computational fabric.
In return, the program must describe more precisely where data resides and how it is distributed.
Who should examine their code?
This change does not imply that every ordinary pmap program will fail.
The public API and many common uses remain recognizable.
The higher-risk cases include code that depends on:
PmapShardingor legacypxlainterfacesdevice_put_shardedordevice_put_replicatedautomatic correction of incompatible input shardings
nested
pmapcallsx[0]-style “unreplication” in multi-host codedirect use of objects returned by
pmap(...).lower(...).compile()tooling or extensions coupled to the removed C++
pmapAPIs
JAX 0.10.0 also removed device_put_sharded and device_put_replicated from the public API.
Their replacements use jax.device_put together with explicit NamedSharding.
Why new code should look directly at shard_map
Because the current pmap implementation already uses shard_map, it may seem reasonable to keep using only the familiar wrapper.
For simple cases, that may remain practical.
But JAX recommends that new or important code migrate directly to shard_map.
Doing so provides explicit control over:
the device mesh
how input dimensions map to mesh axes
whether outputs are sharded or replicated
which mesh axes participate in collective operations
whether input resharding is required
where buffer donation can follow a zero-copy path
pmap remains a familiar interface for common parallel patterns.
shard_map exposes the distributed execution contract underneath it.
The API survived. Its implicit guarantees did not all survive with it.
JAX 0.10.0 did not remove pmap.
It removed the separate C++ infrastructure that existed specifically to implement it and moved the API onto the general jit and shard_map stack.
The function name stayed the same.
The scope of what that function handled implicitly became narrower.
The old implementation knew about special replication patterns, silently adjusted some placements, and provided dedicated fast paths.
The new implementation applies JAX’s broader sharding model more consistently.
That simplifies the system and creates a more composable foundation.
It also means that assumptions hidden inside older code may now have to be written explicitly.
The broader lesson is straightforward:
Keeping an API name does not guarantee that the execution contract beneath it stayed the same.
When reviewing a patch, it is not enough to ask whether a function was removed.
You also have to ask what engine now executes it — and which responsibilities that engine no longer accepts on the user’s behalf.
Related material
JAX 0.10.0 Change Log
Migrating to the new
jax.pmapManual parallelism with
shard_map
This article was written after reviewing JAX’s public change log and official migration documentation.
#JAX #OpenSource #DistributedComputing #CodeAnalysis #SoftwareArchitecture #MachineLearning