jax>=0.8.1
msgpack
optax
orbax-checkpoint
tensorstore
rich>=11.1
typing_extensions>=4.2
PyYAML>=5.4.1
treescope>=0.1.7

[:python_version >= "3.11"]
numpy>=1.23.2

[:python_version >= "3.12"]
numpy>=1.26.0

[dev]
nanobind>=2.5.0
pre-commit>=3.8.0
scikit-build-core[pyproject]>=0.11.0

[docs]
sphinx==6.2.1
sphinx-book-theme
Pygments>=2.6.1
ipykernel
myst_nb
nbstripout
recommonmark
ipython_genutils
sphinx-design
jupytext==1.13.8
dm-haiku>=0.0.14
docutils
matplotlib
scikit-learn
ml_collections
einops
kagglehub>=0.3.3
ipywidgets>=8.1.5

[testing]
clu
einops
gymnasium[atari]
jaxlib
jaxtyping
jraph>=0.0.6dev0
ml-collections
mypy
opencv-python
pytest
pytest-cov
pytest-custom_exit_code
pytest-xdist
pytype
sentencepiece==0.2.0
tensorflow_datasets
torch
cloudpickle>=3.0.0
ale-py>=0.10.2

[testing:platform_system != "Darwin" and python_version < "3.13"]
tensorflow_text>=2.11.0

[testing:python_version < "3.10"]
clu<=0.0.9

[testing:python_version < "3.13"]
tensorflow>=2.12.0

[testing:python_version >= "3.10"]
treescope>=0.1.1

[testing:python_version >= "3.13"]
protobuf<6
tensorflow>=2.20.0
