numpy
scipy
matplotlib
flax
ml_collections
tqdm
absl-py
wandb

[examples]
optax
orbax-checkpoint
torch
mlkernels

[linting]
flake8
pylint
mypy
typing-extensions
pre-commit
ruff
jaxtyping

[testing]
optax
orbax-checkpoint
torch
pytest
pytest-xdist
pytest-cov
coveralls
jax>=0.4.1
jaxlib>=0.4.1
setuptools_scm[toml]
setuptools_scm_git_archive
