hf/transformers/examples/flax/language-modeling/requirements.txt

datasets >= 1.1.3
jax>=0.2.8
jaxlib>=0.1.59
flax>=0.3.5
optax>=0.0.9