search

patrick-kidger / equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Apache License 2.0
2.12k stars 142 forks source link

Issue with Equinox compatibility after `jax==0.4.34` update #869

Closed Qazalbash closed 1 month ago

Qazalbash commented 1 month ago

Hi, I wanted to bring to your attention that the recent release of jax==0.4.34 contains breaking changes that are affecting Equinox. As a result, I am encountering the following error. Additionally, several unit tests are failing due to these changes.

Error

```bash File "/home/muhammad.zeeshan/.conda/envs/kvenv/lib/python3.11/site-packages/flowMC/nfmodel/base.py", line 5, in import equinox as eqx File "/home/muhammad.zeeshan/.conda/envs/kvenv/lib/python3.11/site-packages/equinox/__init__.py", line 3, in from . import debug as debug, internal as internal, nn as nn File "/home/muhammad.zeeshan/.conda/envs/kvenv/lib/python3.11/site-packages/equinox/internal/__init__.py", line 46, in from ._finalise_jaxpr import ( File "/home/muhammad.zeeshan/.conda/envs/kvenv/lib/python3.11/site-packages/equinox/internal/_finalise_jaxpr.py", line 186, in from ._noinline import noinline_p File "/home/muhammad.zeeshan/.conda/envs/kvenv/lib/python3.11/site-packages/equinox/internal/_noinline.py", line 379, in jax.core.pp_eqn_rules[noinline_p] = _noinline_pretty_print ^^^^^^^^^^^^^^^^^^^^^ File "/home/muhammad.zeeshan/.conda/envs/kvenv/lib/python3.11/site-packages/jax/_src/deprecations.py", line 52, in getattr raise AttributeError(message) AttributeError: jax.core.pp_eqn_rules was removed in JAX v0.4.34. ```

Failing testcases

```bash ============================================================= short test summary info ============================================================= SKIPPED [1] tests/test_onnx.py:8: unconditional skip FAILED tests/test_scan.py::test_scan[True-4-checkpointed] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_scan.py::test_scan[True-None-checkpointed] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_forward[checkpointed-False] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_forward[checkpointed-True] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-1-3-0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-1-5-0,1,2,3,4,0,1,2,3,0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-2-8-7,0,1,2,3,4,5,6,3,4,5,3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-3-11-10,4,5,6,7,8,9,7,8,7,4,5,6,5,4,0,1,2,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-9-5-4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-2-5-3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-4-3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-5-3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-6-5,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-7-6,5,2,3,4,3,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-8-7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-9-7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-10-9,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-11-10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-12-10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-13-12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-True-4-14-12,13,12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-1-3-0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-1-5-0,1,2,3,4,0,1,2,3,0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-2-8-7,0,1,2,3,4,5,6,3,4,5,3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-3-11-10,4,5,6,7,8,9,7,8,7,4,5,6,5,4,0,1,2,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-9-5-4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-2-5-3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-4-3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-5-3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-6-5,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-7-6,5,2,3,4,3,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-8-7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-9-7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-10-9,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-11-10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-12-10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-13-12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[False-False-4-14-12,13,12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-1-3-0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-1-5-0,1,2,3,4,0,1,2,3,0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-2-8-7,0,1,2,3,4,5,6,3,4,5,3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-3-11-10,4,5,6,7,8,9,7,8,7,4,5,6,5,4,0,1,2,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-9-5-4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-2-5-3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-4-3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-5-3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-6-5,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-7-6,5,2,3,4,3,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-8-7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-9-7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-10-9,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-11-10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-12-10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-13-12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-True-4-14-12,13,12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-1-3-0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-1-5-0,1,2,3,4,0,1,2,3,0,1,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-2-8-7,0,1,2,3,4,5,6,3,4,5,3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-3-11-10,4,5,6,7,8,9,7,8,7,4,5,6,5,4,0,1,2,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-9-5-4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-2-5-3,4,3,0,1,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-4-3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-5-3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-6-5,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-7-6,5,2,3,4,3,2,0,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-8-7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-9-7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-10-9,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-11-10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-12-10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-13-12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_backward_checkpointed[True-False-4-14-12,13,12,9,10,11,10,9,5,6,7,8,7,6,5,0,1,2,3,4,3,2,1,0] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_vmap_primal_unbatched_cond[checkpointed-False] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_vmap_primal_unbatched_cond[checkpointed-True] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_vmap_primal_batched_cond[checkpointed-False] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_vmap_primal_batched_cond[checkpointed-True] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_vmap_cotangent[checkpointed-False] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_vmap_cotangent[checkpointed-True] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_speed_buffer_while[False] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_speed_buffer_while[True] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_speed_grad_checkpointed_while - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_nested_loops[False] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_nested_loops[True] - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_zero_buffer - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_symbolic_zero - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_nondifferentiable_body1 - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_nondifferentiable_body2 - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_body_fun_grads - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_trivial_vjp - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: FAILED tests/test_while_loop.py::test_unperturbed_output - TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got: ============================================== 90 failed, 362 passed, 1 skipped in 197.63s (0:03:17) ============================================== ```

patrick-kidger commented 1 month ago

Upgrade to a more modern version of Equinox. This has already been fixed in recent versions. :)

Qazalbash commented 1 month ago

I rechecked it on the main branch of this repo, they are failing.

$ pip list | grep -E '(jax|equinox|numpy)'                                                
equinox           0.11.7
jax               0.4.34
jaxlib            0.4.34
jaxtyping         0.2.34
numpy             2.1.2
patrick-kidger commented 1 month ago

I really don't think so. I don't know how you're getting version 0.11.7 reported, but take a look at the relevant line of the source code in this repository -- the offending jax.core.pp_eqn_rules is no longer present. And just to be sure I've double-checked this myself on both HEAD and the 0.11.7 release.

Qazalbash commented 1 month ago

I cleaned the caches and run again, now I see it has been resolved. Sorry for the trouble.