absl-py>=0.12.0
immutabledict>=2.2.1
numpy>=1.19.5
distrax>=0.1.2
chex>=0.1.3
jax>=0.3.13
jaxlib>=0.3.10

[tests]
pytest-xdist
absl-py==0.12.0
immutabledict==2.2.1
numpy==1.19.5
distrax==0.1.2
chex==0.1.3
jax==0.3.13
jaxlib==0.3.10
dm-haiku==0.0.6
