jaxlib
jax
jaxopt
numpy
flax
tqdm
wget
h5py
imt-tree-utils
wandb
optax
dm-haiku
joblib
pyyaml
qmt
scipy

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