Pmap — это функция JAX, которая запускает копии функции на нескольких аппаратных устройствах, таких как CPU, GPU или TPU. Она поддерживает выполнение single-program multiple-data, применяя функцию к разным фрагментам входных данных.
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.