 code, each on different data inputs. More sophisticated approaches are possible, and you can go over to the JAX documentation to explore this topic in more detail if you are interested! jax.pmap does two things: Compiles (or ”jits”) the code which is similar to jax.jit(). This does not happen when you call pmap, and only the first time the pmapped function is called. Ensures the compiled code runs in parallel on all available devices. To demon