jax
optax
matplotlib

[dev]
torch
torchvision
