chex~=0.1.88
jax~=0.4.38
tensorflow-probability[jax]~=0.25.0

[graph]
jraph~=0.0.6.dev0
