jax
jaxlib
matplotlib
numpy
scipy
tensorflow_probability
tqdm
dm-haiku
optax
jaxopt
