diffrax
einops
equinox
jaxtyping
matplotlib
ml-dtypes
numpy
optax
optimistix
optuna
plum-dispatch
torch
torchvision
tqdm
wadler_lindig
wandb
pandas
scikit-learn
seaborn