tensorflow-probability
numpyro
jax<=0.2.20

[test]
pytest>=3.6
