jax
numpy

[tests]
absl-py
jaxlib
numpy>=1.17
pytest
