clu
datasets
distrax
grain
matplotlib
seaborn
tensorflow
tensorflow-datasets
tf-keras
transformers
# for jax on TPU
--find-links https://storage.googleapis.com/jax-releases/jax_releases.html
--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
jax[tpu]
