Glossary term
Glossary term
Infrastructure and Serving
A JAX function that executes copies of an input function on multiple underlying hardware devices (CPUs, GPUs, or TPUs), with different input values. pmap relies on SPMD.
Created for this library
A research engineer uses pmap in JAX to apply the same training step across many devices with minimal code change.
An ML platform team uses pmap to scale gradient computation across multiple accelerators in JAX-based training.
A research team uses pmap as the simplest path to multi-device data parallelism in JAX during prototyping.
Definition source: Google for Developers Machine Learning Glossary | Creative Commons Attribution 4.0 License