jax
jaxlib
optax
numpy
absl-py
pytest
matplotlib
