transformers
torch
jax==0.4.23
jaxlib==0.4.23
flax==0.7.5
optax==0.1.7