numpy
jax>=0.1.73
jaxlib>=0.1.52
dataclasses

[docs]
sphinx>=3
sphinx_rtd_theme
sphinx-autodoc-typehints==1.11.1
myst-nb
numpyro
matplotlib
h5py
tensorflow
tensorflow_datasets

[tensorflow]
tensorflow
tensorflow_datasets

[test]
pylint
pytest
pytest-mock
