ite you to go over the JAX documentation and the pjit pages to explore this topic if you are interested!
jax.pmap does two things for us:
Compiles (or jits) the code, as if we had invoked jax.jit(). This does not happen when we call pmap, but the first time the pmapped function is invoked.
Ensures the compiled code runs in parallel in all the available devices.
To show how it works we pmap the _generate method of the pipeline, which is the priv