jax
jaxlib

[testing]
flax
pytest
pytest-cov
