jax>=0.4.38
jax-tqdm>=0.4.0

[notebooks]
flax
jupyter
matplotlib
notebook
rejax>=0.1.2
tqdm
