flax
nltk
optax
tensorflow-datasets