flax
jax
huggingface_hub
optax
numpy
tiktoken
flash_attention_jax
