datasets
flax>=0.11.1
fsspec
google-metrax>=0.2.3
grain
huggingface_hub
jaxtyping
jinja2
kagglehub
numba
omegaconf
pylatexenc
python-dotenv
qwix
sentencepiece
sympy
tensorflow_datasets
tqdm
transformers
hf_transfer

[dev]

[docs]
sphinx>=8.2.3
sphinx-book-theme>=1.1.4
sphinx-autodoc-typehints
ipython>=8.8.0
myst-nb>=1.3.0
matplotlib>=3.10.0
sphinx-gallery>=0.19.0
sphinx-collections>=0.0.1
sphinx_contributors

[prod]
jax[tpu]!=0.7.2,>=0.6.0
