transformers
torch
jaxjax==0.4.13
jaxlibjax==0.4.13
flax
optax