chex
datasets==2.14.3
distrax
einops>=0.6.1
fjformer>=0.0.5
flax
gradio~=3.41.2
ipython
jax>=0.4.10
jaxlib>=0.4.10
msgpack
optax
pydantic==2.4.2
rlax
tensorboard
tqdm
transformers>=4.33.0
typing
wandb>=0.15.9
