While running experiments, I also often get this warning which seems to say that it might be problematic:
UserWarning: The jitted function <unnamed function> includes a pmap. Using jit-of-pmap can lead to inefficient data movement, as the outer jit does not preserve sharded data representations and instead collects input and output arrays onto a single device. Consider removing the outer jit unless you know what you're doing. See [https://github.com/google/jax/issues/2926].
Hi, thanks for raising this.
I'm aware of this warning and it is indeed intended.
If you run experiments on multiple devices (e.g. this, which by defaults uses colab TPUs), you will notice the difference with and without this outer jit.
https://github.com/google/evojax/blob/f41a6e01cbfb9ace1bd967e86460132dd210e5d1/evojax/sim_mgr.py#L189
The docs seem to say that jits around pmaps are unnecessary: https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html#pmap-and-jit
While running experiments, I also often get this warning which seems to say that it might be problematic:
UserWarning: The jitted function <unnamed function> includes a pmap. Using jit-of-pmap can lead to inefficient data movement, as the outer jit does not preserve sharded data representations and instead collects input and output arrays onto a single device. Consider removing the outer jit unless you know what you're doing. See [https://github.com/google/jax/issues/2926].
If the behaviour is intended please discard.