jax>=0.4.20

[dev]
pytest>=7.0
numpy>=1.24
