absl-py>=0.12.0
immutabledict>=2.2.1
numpy>=1.22
distrax>=0.1.3
jax>=0.4.27
jaxlib>=0.4.27
dm-tree>=0.1.7
optax>=0.1.4
typing-extensions>=4.0.0

[docs]
sphinx>=5.3.0
sphinx_rtd_theme>=1.3.0
sphinxcontrib-katex>=0.9.7
sphinxcontrib-bibtex>=2.6.1
sphinx-autodoc-typehints>=1.20.0
pandoc>=2.3
IPython>=8.10.0
ipykernel>=5.3.4
myst_nb>=0.17.2
docutils>=0.17.0
matplotlib>=3.5.0

[examples]
ml_collections>=0.1.1
optax>=0.1.1
dm-haiku>=0.0.6
jaxline>=0.0.5
tensorflow>=2.8.0
tensorflow_datasets>=4.5.2

[tests]
pytest-xdist
absl-py>=0.12.0
immutabledict>=2.2.1
numpy>=1.22
distrax>=0.1.3
jax>=0.6.0
jaxlib>=0.4.27
dm-haiku>=0.0.9
dm-tree>=0.1.7
optax>=0.1.4
