absl-py>=0.12.0
immutabledict>=2.2.1
numpy>=1.21
distrax>=0.1.3
jax>=0.4.7
jaxlib>=0.4.7
dm-tree>=0.1.7
optax>=0.1.4

[:python_version < "3.10"]
typing_extensions>=4.2.0

[tests]
pytest-xdist
absl-py==0.12.0
immutabledict==2.2.1
numpy==1.21
distrax==0.1.3
jax==0.4.7
jaxlib==0.4.7
dm-haiku==0.0.9
dm-tree==0.1.7
optax==0.1.4
