Bounded versions of JAX and diffrax due to JAX bug

[DanPuzzuoli]

The change in #189 should be removed when the corresponding JAX bug discussed here is fixed.

Edit:

• Now both PRs Temporary JAX fix for tests #189 and Temporary upper bound on JAX version <=0.4.6 #210 need to be reverted once fixed.
• For reference, the failing tests are the methods test_jit_grad_dyson_magnus and test_jit_grad_dyson_like of test.dynamics.perturbation.test_solve_lmde_perturbation.Testsolve_lmde_perturbationJAX.

[DanPuzzuoli]

As of JAX 0.4.7 the change in #189 can no longer be used to bypass the issue with JAX. For now we may need to set 0.4.6 as an upper bound on the JAX version.

[DanPuzzuoli]

In relation to the previous comment #210 now sets 0.4.6 to the latest version of JAX that works with Dynamics. For now this is just to get the CI tests working - hopefully this will be resolved shortly and these restrictions can be removed.

[DanPuzzuoli]

As this issue has continued to evolve, I'll make this comment a summary of what's going on, to be edited as new issues arise.

JAX 0.4.4 introduced a new way of tracing functions. This new way unfortunately has a bug that breaks the perturbation module, but until JAX 0.4.7, it was possible to revert back to the old way. As this bug is still not fixed, it is currently necessary to bound the JAX version to 0.4.6. The latest version of diffrax now also assumes the latest version of JAX, and hence we need to bound diffrax and equinox to suitable versions. Hopefully this bug gets fixed soon and we can go back to using the latest version of all of these packages.

List of issues/PRs:

• Temporary JAX fix for tests #189 sets the test environment to revert to the old way of doing tracing.
• Temporary upper bound on JAX version <=0.4.6 #210 bounds the JAX version in setup.py.
• Bound Diffrax version #226 bounds the diffrax and equinox versions in setup.py.
• Add warning for JAX versions on import of Dynamics #232 inserts a warning in the Array dispatch setup if a JAX version >=0.4.4 is being used. (This warning is triggered when Dynamics is imported.) docs/conf.py now sets os.environ["JAX_JIT_PJIT_API_MERGE"] = "0" for the docs build to avoid this warning.

Note: the issue in the perturbation module arises due to a weird combination of JAX elements outlined here. It would be worth checking if this problem occurs if we sub a diffrax solver in for jax_odeint. My guess is it won't make a difference, but it's worth a shot.