jax
tqdm
optax
jaxtyping
rich

[:python_version < "3.11"]
typing-extensions

[test]
pytest
pytest-cov
pytest-mock
chex
