jaxlib
optax
pytest
beartype
