@murphyk Currently, we pin jax/jaxlib versions in each release. From the next release, we’ll relax that restriction. For now, the simplest way is to add
!pip install -q numpyro@git+https://github.com/pyro-ppl/numpyro
to the top of your colab notebook, like bayesian regression tutorial