datasets
flax>=0.7.1
nltk
optax