jax>=0.6.0
jaxlib
matplotlib
numpy
scipy
tensorflow_probability

[examples]
scikit-learn
optax

[experimental]

[tests]
scikit-learn
networkx
psutil
pytest
flake8
