chex~=0.1.84
datasets~=2.14.6
distrax
einops~=0.6.1
fastapi~=0.104.1
fjformer>=0.0.20
flax~=0.7.5
gradio~=4.12.0
ipython~=8.17.2
jax>=0.4.17
jaxlib>=0.4.17
msgpack~=1.0.7
numpy~=1.26.2
optax~=0.1.7
pydantic-core==2.14.6
pydantic==2.5.3
rlax
setuptools~=68.1.2
tensorboard
torch~=2.1.0
tqdm~=4.64.1
transformers>=4.33.0
typing~=3.7.4.3
uvicorn~=0.23.2
wandb>=0.15.9
