transformers
torch
jaxjax==0.4.23
jaxlibjax==0.4.23
flax==0.7.5
optax==0.1.7