flax
optax
einops
tensorflow
tensorflow-datasets
git+https://github.com/deepmind/dm-haiku
torch
torchvision
pandas
