numpy
jax
optax
equinox
jax-tqdm
torch

[notebook]
jupyter
matplotlib
