Probabilistic Machine Learning in NumPyro and Jax