jaxlib
jax
jaxopt
numpy
flax
tqdm
imt-tree-utils
wandb
optax
dm-haiku
pyyaml

[dev]
mkdocs
mkdocs-material
mkdocstrings
mkdocstrings-python
mknotebooks
pytest
pytest-xdist
nbmake
