jax
jaxlib
numpyro
tensorflow-probability
inference-gym
