flax
jax
jaxlib

[testing]
pytest
