MrVI slowdown due to JAX compilation update

[justjhong]

With recent updates to JAX, MrVI trains significantly slower than before. We suspect it is due to the new AOT compilation strategy (https://jax.readthedocs.io/en/latest/aot.html).

Any basic training with MrVI with a fresh install. Reproduced by @PierreBoyeau and myself.

[ori-kron-wis]

@justjhong lmk if you want me to upper bound jax for now and to which version.
We will probably have a release in the coming week or so, so it can get into there.

[justjhong]

Hi @ori-kron-wis, thanks for checking. I took some time this morning to try to debug it but was not able to find a solution.
I was able to figure out that the problems arise starting from jax==0.4.36. So for now, let's upper bound to jax<0.4.36 (non-inclusive).

[PierreBoyeau]

Hi, here are a few more details about my understanding of the problem.
It seems that lightning introduces a significant overhead when training the model, here and there for instance.
MrVI is notably faster at this stage without the lightning wrapper.

[canergen]

For now pinning jax<0.4.36. Potentially related to jax-ml/jax#26162. Check again when this is adressed. Leaving this open as pinning circumvents it but might create issues in the near future.