jax
equinox
pydantic
numpy
jaxlib
tqdm
beartype
typing_extensions
loguru
optax
jaxtyping
pytest
ftfy
regex
statedict2pytree>=1.0.2
jaxonlayers>=0.1.4

[dev]
pre-commit
pytest
mkdocs

[examples]
tensorflow-cpu
tensorflow_datasets
clu
requests

[tests]
torch
torchvision
