jax
jaxlib
optax
numpy
absl-py

[test]
pytest

[viz]
matplotlib
