numpy>=1.22
jax>=0.4.27
msgpack
optax
orbax-checkpoint
tensorstore
rich>=11.1
typing_extensions>=4.2
PyYAML>=5.4.1

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

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

[all]
matplotlib

[testing]
clu
einops
gymnasium[accept-rom-license,atari]
jaxlib
jaxtyping
jraph>=0.0.6dev0
ml-collections
mypy
opencv-python
pytest
pytest-cov
pytest-custom_exit_code
pytest-xdist
pytype
sentencepiece
tensorflow_text>=2.11.0
tensorflow_datasets
tensorflow>=2.12.0
torch
nbstripout
black[jupyter]==23.7.0

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

[testing:python_version >= "3.10"]
penzai>=0.1.2
