jax
numpy

[tests]
absl-py
chex
pytest
