attrs
jax
jaxtyping
equinox
optax
dpshdl
chex
importlib-resources
cloudpickle
pillow
omegaconf
gitpython
tensorboard
psutil
requests

[all]
black
darglint
mypy
ruff
pytest
types-pillow
types-psutil
types-requests
orbax-export
tensorflow
flax

[dev]
black
darglint
mypy
ruff
pytest
types-pillow
types-psutil
types-requests

[export]
orbax-export
tensorflow

[flax]
flax
