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
typing_extensions>=4.2.0; python_version<"3.10"
