jax~=0.8.0
tfp-nightly[jax]

[graph]
jraph~=0.0.6.dev0
