jax[cpu]

[dev]
scikit-learn
pandas
nbdev
jupyter
dm-haiku
optax
torch
torchvision
datasets
