flax
transformers
jax
jaxlib
numpy
torch