jax[tpu]>=0.2.16
flax
optax
datasets
datasets[streaming]
jsonlines 
zstandard
transformers
tokenizers
sentencepiece
wandb
einops
nltk