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