pytest-xdist
absl-py==0.12.0
immutabledict==2.2.1
numpy==1.20
distrax==0.1.2
chex==0.1.3
jax==0.3.17
jaxlib==0.3.15
dm-haiku==0.0.7
dm-tree==0.1.7
