flax>=0.8.5
jax>=0.4.31
optax
numpy
tqdm
mergedeep
typing
orbax-checkpoint>=0.5.15
