matplotlib
numpy
jax
optax
equinox
imageio
h5py
