flax
jax
optax
numpy
tqdm
mergedeep
typing
orbax-checkpoint
