[Question] Are the jits around pmaps intended?

[MaximilienLC]

evojax/evojax/sim_mgr.py

Line 189 in f41a6e0

self._train_rollout_fn = jax.jit(jax.pmap(

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.

[lerrytang]

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.