datasets
flax
optax
Pillow