jax
jaxlib
equinox
optax
zodiax
matplotlib
jupyter
jupyterlab
tqdm
chainconsumer
