jax
optax
tqdm
