flax>=0.8.4
jax>=0.4.28
optax>=0.2.2
orbax
numpy
clu
einops
tqdm
grain
termcolor
augmax
albumentations
rich
python-dotenv
