jax>=0.4.0
jaxlib>=0.4.0
jaxopt
optax
numpy
matplotlib
seaborn
