chex~=0.1.85
jax~=0.4.23
jaxlib~=0.4.23
matplotlib~=3.8.2
tensorflow-probability[jax]~=0.23.0

[graph]
jraph~=0.0.6.dev0
