Speeding up prototyping when using pmap

[dedeswim]

Hi all!

First of all, thanks a lot for creating JAX, I am truly enjoying working with it.

I am lucky to have access to some TPUs thanks to the TRC program, and for this reason I am using pmap inside my codebase. However, in some cases, instead of doing some long-running tasks, I would simply like to run 1 step of my training loop to test if my changes work as expected. Unfortunately, using pmap, my training loop always JIT compiles making this process very slow (the compilation time is ~288s, as the training step involves adversarial training, which I guess generates a lot of code).

Is there some flag I can use to quickly prevent my script to use pmap, and instead run my function without compiling it? (i.e. something similar to the JAX_DISABLE_JIT flag, but for pmap)

Thanks!