jax
jaxlib
typing_extensions

[testing]
flax
pytest
pytest-cov
