tqdm
optax
matplotlib
jax
einops
scikit-learn
equinox
diffrax
ott-jax
optimistix
