flax
jax
jaxlib
optax
einops
sentencepiece
