equinox >= 0.11.0, <= 0.11.11
jax >= 0.3.2, <= 0.5.0
numpy >= 1.20.0, <= 2.2.2
