jax
jaxlib
jax-verify
optax
matplotlib
numpy
