jax
optax
numpy
plotly
matplotlib
tqdm
