Pmap

Parallel map over devices.

Abstract Signature:

Pmap(f: Callable)

PyTorchJAX (Core)Flax NNX

PyTorch

API: —

Strategy: Custom / Partial

JAX (Core)

API: jax.pmap

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.pmap

Strategy: Direct Mapping

Official Docs ↗