absl-py>=0.12.0
immutabledict>=2.2.1
numpy>=1.19.5
distrax>=0.1.2
chex>=0.1.3
jax>=0.3.17
jaxlib>=0.3.15
dm-tree>=0.1.7
