Pmap is a JAX function that runs copies of a function across multiple hardware devices such as CPUs, GPUs, or TPUs. It supports single-program multiple-dataexecution by applying the function to different input shards.
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.