jax.pmap ======== .. currentmodule:: jax .. autofunction:: pmap