jax
jax_tqdm
matplotlib