jax>=0.6.0
jaxctx>=1.2.0
matplotlib
numpy
scipy
tfp-nightly

[:python_version < "3.11"]
tomli

[distributed]
cloudpickle>=3
pyzmq>=25

[examples]
scikit-learn
optax

[tests]
scikit-learn
networkx
psutil
pytest
flake8
ruff
cloudpickle>=3
pyzmq>=25
