jax>=0.5.0
jax-tqdm
numpy
scipy
tqdm
