Pin jax version temporarily

[dpanici]

Pin jax to <0.4.34 until we fix bug listed in #1291

[rahulgaur104]

Merge if all the tests pass.

[f0uriest]

Would it be easier to just fix the underlying issue?

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.60 +/- 8.25     | +3.17e-03 +/- 4.34e-02 |  5.30e-01 +/- 2.4e-02  |  5.26e-01 +/- 3.6e-02  |
 test_equilibrium_init_medres            |     +3.05 +/- 5.43     | +1.26e-01 +/- 2.25e-01 |  4.27e+00 +/- 2.2e-01  |  4.14e+00 +/- 3.0e-02  |
 test_equilibrium_init_highres           |     +0.82 +/- 1.04     | +4.51e-02 +/- 5.71e-02 |  5.52e+00 +/- 3.6e-02  |  5.47e+00 +/- 4.5e-02  |
 test_objective_compile_dshape_current   |     +4.49 +/- 5.27     | +1.67e-01 +/- 1.96e-01 |  3.89e+00 +/- 4.6e-02  |  3.72e+00 +/- 1.9e-01  |
 test_objective_compute_dshape_current   |     +1.56 +/- 1.76     | +5.62e-05 +/- 6.36e-05 |  3.67e-03 +/- 4.3e-05  |  3.61e-03 +/- 4.7e-05  |
 test_objective_jac_dshape_current       |     -1.26 +/- 7.98     | -5.11e-04 +/- 3.23e-03 |  4.00e-02 +/- 2.3e-03  |  4.05e-02 +/- 2.3e-03  |
 test_perturb_2                          |     +0.49 +/- 1.19     | +8.59e-02 +/- 2.09e-01 |  1.77e+01 +/- 1.5e-01  |  1.76e+01 +/- 1.5e-01  |
 test_proximal_freeb_jac                 |     -0.01 +/- 1.20     | -6.51e-04 +/- 8.99e-02 |  7.51e+00 +/- 5.2e-02  |  7.51e+00 +/- 7.4e-02  |
 test_solve_fixed_iter                   |     +1.68 +/- 61.14    | +8.38e-02 +/- 3.04e+00 |  5.06e+00 +/- 2.1e+00  |  4.97e+00 +/- 2.2e+00  |
 test_build_transform_fft_midres         |     -0.14 +/- 4.92     | -8.38e-04 +/- 3.03e-02 |  6.15e-01 +/- 1.4e-02  |  6.16e-01 +/- 2.7e-02  |
 test_build_transform_fft_highres        |     +0.99 +/- 3.11     | +1.00e-02 +/- 3.13e-02 |  1.02e+00 +/- 2.0e-02  |  1.01e+00 +/- 2.4e-02  |
 test_equilibrium_init_lowres            |     -0.72 +/- 2.60     | -2.77e-02 +/- 1.00e-01 |  3.82e+00 +/- 4.6e-02  |  3.85e+00 +/- 8.9e-02  |
 test_objective_compile_atf              |     +0.17 +/- 3.74     | +1.29e-02 +/- 2.93e-01 |  7.85e+00 +/- 2.0e-01  |  7.84e+00 +/- 2.1e-01  |
 test_objective_compute_atf              |     +1.36 +/- 2.03     | +1.42e-04 +/- 2.13e-04 |  1.06e-02 +/- 1.4e-04  |  1.05e-02 +/- 1.6e-04  |
 test_objective_jac_atf                  |     -0.60 +/- 1.93     | -1.19e-02 +/- 3.80e-02 |  1.96e+00 +/- 2.2e-02  |  1.97e+00 +/- 3.1e-02  |
 test_perturb_1                          |     -0.08 +/- 1.86     | -9.87e-03 +/- 2.36e-01 |  1.27e+01 +/- 1.8e-01  |  1.27e+01 +/- 1.5e-01  |
 test_proximal_jac_atf                   |     -0.43 +/- 0.76     | -3.49e-02 +/- 6.21e-02 |  8.16e+00 +/- 4.4e-02  |  8.20e+00 +/- 4.4e-02  |
 test_proximal_freeb_compute             |     +0.26 +/- 1.35     | +4.83e-04 +/- 2.50e-03 |  1.86e-01 +/- 1.9e-03  |  1.85e-01 +/- 1.6e-03  |

[rahulgaur104]

If it's an easy fix, we can always make a new PR, and bump up the dependency.
If not, the issue can be fixed later.
This PR doesn't affect anything negatively.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.45%. Comparing base (e35436d) to head (e506d80).
Report is 2 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1292   +/-   ##
=======================================
  Coverage   95.45%   95.45%           
=======================================
  Files          96       96           
  Lines       23806    23806           
=======================================
+ Hits        22723    22725    +2     
+ Misses       1083     1081    -2     

see 1 file with indirect coverage changes

[dpanici]

Would it be easier to just fix the underlying issue?

Yes but I am working to finish my APS poster right now... so this is a temporary fix

[YigitElma]

There are some other bugs related to JAX 0.4.34, for example equinox stuff, AttributeError: jax.core.pp_eqn_rules was removed in JAX v0.4.34.