jaxlib
jax>=0.2.26
tqdm
