setuptools
flashbax
equinox
optax
mctx
numpy
scipy
wandb
matplotlib
tqdm
distrax