jax
tqdm
jaxtyping
equinox>=0.13.2
numpy

[test]
pytest
pytest-cov
chex
