jax
flax
numpy
matplotlib
scikit-image
natsort
sympy
k3d
absl-py
seaborn
optax
h5py
equinox
tqdm
wandb
pyyaml
ml-collections
diffrax
jaxlib
orbax-checkpoint
git+https://github.com/google-deepmind/kfac-jax
git+https://github.com/haydn-jones/SOAP_JAX
