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