Repository: pyro-ppl/numpyro (language: python)

Issue:
Inference for `gaussian_hmm` is broken on latest jax version (0.4.30)
The following works with jax/jaxlib 0.4.29, but is broken with jax/jaxlib 0.4.30:

Copying from the [current documentation](https://num.pyro.ai/en/stable/primitives.html#numpyro.contrib.control_flow.scan):

```python
import numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.contrib.control_flow import scan

def gaussian_hmm(y=None, T=10):
    def transition(x_prev, y_curr):
        x_curr = numpyro.sample('x', dist.Normal(x_prev, 1))
        y_curr = numpyro.sample('y', dist.Normal(x_curr, 1), obs=y_curr)
        return x_curr, (x_curr, y_curr)

    x0 = numpyro.sample('x_0', dist.Normal(0, 1))
    _, (x, y) = scan(transition, x0, y, length=T)
    return (x, y)


with numpyro.handlers.seed(rng_seed=0):  # generative
    x, y = gaussian_hmm()
```

If I then try to perform inference:
```python
from jax import random
from numpyro.infer import MCMC, NUTS
nuts_kernel = NUTS(gaussian_hmm)

mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000)
rng_key = random.PRNGKey(0)

mcmc.run(rng_key, y=y)
```

I get the error
```pycon
TypeError: body_fun output and input must have identical types, got
('ShapedArray(int32[], weak_type=True)', ['ShapedArray(float32[10])', 'DIFFERENT ShapedArray(int32[], weak_type=True) vs. ShapedArray(float0[])', 'ShapedArray(float32[])'], []).
```


Discussion / hints:
Hi @damonbayer, it is fixed in #1817. We will make a patch release this week.

Fix patch:
diff --git a/examples/annotation.py b/examples/annotation.py
index 3341dbe72..a28ab48be 100644
--- a/examples/annotation.py
+++ b/examples/annotation.py
@@ -320,7 +320,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Bayesian Models of Annotation")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/ar2.py b/examples/ar2.py
index b6f0e737a..6ebb5949e 100644
--- a/examples/ar2.py
+++ b/examples/ar2.py
@@ -138,7 +138,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="AR2 example")
     parser.add_argument("--num-data", nargs="?", default=142, type=int)
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
diff --git a/examples/baseball.py b/examples/baseball.py
index 1588d5898..47615c37a 100644
--- a/examples/baseball.py
+++ b/examples/baseball.py
@@ -210,7 +210,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Baseball batting average using MCMC")
     parser.add_argument("-n", "--num-samples", nargs="?", default=3000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1500, type=int)
diff --git a/examples/bnn.py b/examples/bnn.py
index 9dd29d291..022e88d9e 100644
--- a/examples/bnn.py
+++ b/examples/bnn.py
@@ -160,7 +160,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Bayesian neural network example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=2000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/covtype.py b/examples/covtype.py
index 8b59b99dc..424eac7f7 100644
--- a/examples/covtype.py
+++ b/examples/covtype.py
@@ -206,7 +206,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="parse args")
     parser.add_argument(
         "-n", "--num-samples", default=1000, type=int, help="number of samples"
diff --git a/examples/funnel.py b/examples/funnel.py
index 2e9c3d1b6..bb28e7dfa 100644
--- a/examples/funnel.py
+++ b/examples/funnel.py
@@ -139,7 +139,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(
         description="Non-centered reparameterization example"
     )
diff --git a/examples/gaussian_shells.py b/examples/gaussian_shells.py
index 98870b167..9c2184eb4 100644
--- a/examples/gaussian_shells.py
+++ b/examples/gaussian_shells.py
@@ -120,7 +120,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Nested sampler for Gaussian shells")
     parser.add_argument("-n", "--num-samples", nargs="?", default=10000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/gp.py b/examples/gp.py
index 0f70400db..dfa22f624 100644
--- a/examples/gp.py
+++ b/examples/gp.py
@@ -180,7 +180,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Gaussian Process example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/hmm.py b/examples/hmm.py
index c414ba59d..3d17b9366 100644
--- a/examples/hmm.py
+++ b/examples/hmm.py
@@ -263,7 +263,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Semi-supervised Hidden Markov Model")
     parser.add_argument("--num-categories", default=3, type=int)
     parser.add_argument("--num-words", default=10, type=int)
diff --git a/examples/holt_winters.py b/examples/holt_winters.py
index 621f0eb35..fb97e62bd 100644
--- a/examples/holt_winters.py
+++ b/examples/holt_winters.py
@@ -180,7 +180,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Holt-Winters")
     parser.add_argument("--T", nargs="?", default=6, type=int)
     parser.add_argument("--future", nargs="?", default=1, type=int)
diff --git a/examples/horseshoe_regression.py b/examples/horseshoe_regression.py
index 834483aa6..ee1624cc2 100644
--- a/examples/horseshoe_regression.py
+++ b/examples/horseshoe_regression.py
@@ -162,7 +162,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Horseshoe regression example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=2000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/minipyro.py b/examples/minipyro.py
index aeea67d0d..8d7094de3 100644
--- a/examples/minipyro.py
+++ b/examples/minipyro.py
@@ -58,7 +58,7 @@ def body_fn(i, val):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Mini Pyro demo")
     parser.add_argument("-f", "--full-pyro", action="store_true", default=False)
     parser.add_argument("-n", "--num-steps", default=1001, type=int)
diff --git a/examples/mortality.py b/examples/mortality.py
index 27a624511..0ee18dc06 100644
--- a/examples/mortality.py
+++ b/examples/mortality.py
@@ -220,7 +220,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
 
     parser = argparse.ArgumentParser(description="Mortality regression model")
     parser.add_argument("-n", "--num-samples", nargs="?", default=500, type=int)
diff --git a/examples/neutra.py b/examples/neutra.py
index 780e420f2..c47d8f94e 100644
--- a/examples/neutra.py
+++ b/examples/neutra.py
@@ -197,7 +197,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="NeuTra HMC")
     parser.add_argument("-n", "--num-samples", nargs="?", default=4000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/ode.py b/examples/ode.py
index cd5a6eabe..3845dcd9c 100644
--- a/examples/ode.py
+++ b/examples/ode.py
@@ -117,7 +117,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Predator-Prey Model")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
diff --git a/examples/prodlda.py b/examples/prodlda.py
index 1c82b27a7..0a30d202d 100644
--- a/examples/prodlda.py
+++ b/examples/prodlda.py
@@ -315,7 +315,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(
         description="Probabilistic topic modelling with Flax and Haiku"
     )
diff --git a/examples/sparse_regression.py b/examples/sparse_regression.py
index 84f1ae8b5..f0c9781d4 100644
--- a/examples/sparse_regression.py
+++ b/examples/sparse_regression.py
@@ -384,7 +384,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Gaussian Process example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=500, type=int)
diff --git a/examples/stochastic_volatility.py b/examples/stochastic_volatility.py
index 7f1583b94..a4ce6d7fd 100644
--- a/examples/stochastic_volatility.py
+++ b/examples/stochastic_volatility.py
@@ -122,7 +122,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Stochastic Volatility Model")
     parser.add_argument("-n", "--num-samples", nargs="?", default=600, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=600, type=int)
diff --git a/examples/thompson_sampling.py b/examples/thompson_sampling.py
index 1e61e4890..e10205b72 100644
--- a/examples/thompson_sampling.py
+++ b/examples/thompson_sampling.py
@@ -292,7 +292,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Thompson sampling example")
     parser.add_argument(
         "--num-random", nargs="?", default=2, type=int, help="number of random draws"
diff --git a/examples/toy_mixture_model_discrete_enumeration.py b/examples/toy_mixture_model_discrete_enumeration.py
index ae68d170a..f64bdb98b 100644
--- a/examples/toy_mixture_model_discrete_enumeration.py
+++ b/examples/toy_mixture_model_discrete_enumeration.py
@@ -126,7 +126,7 @@ def get_true_pred_CPDs(CPD, posterior_param):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Toy mixture model")
     parser.add_argument("-n", "--num-steps", default=4000, type=int)
     parser.add_argument("-o", "--num-obs", default=10000, type=int)
diff --git a/examples/ucbadmit.py b/examples/ucbadmit.py
index f35cb03f7..4147067c6 100644
--- a/examples/ucbadmit.py
+++ b/examples/ucbadmit.py
@@ -151,7 +151,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(
         description="UCBadmit gender discrimination using HMC"
     )
diff --git a/examples/vae.py b/examples/vae.py
index f9f5732d0..5543c30d0 100644
--- a/examples/vae.py
+++ b/examples/vae.py
@@ -160,7 +160,7 @@ def reconstruct_img(epoch, rng_key):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="parse args")
     parser.add_argument(
         "-n", "--num-epochs", default=15, type=int, help="number of training epochs"
diff --git a/notebooks/source/bad_posterior_geometry.ipynb b/notebooks/source/bad_posterior_geometry.ipynb
index 3d61fd14e..68c2f0c56 100644
--- a/notebooks/source/bad_posterior_geometry.ipynb
+++ b/notebooks/source/bad_posterior_geometry.ipynb
@@ -49,7 +49,7 @@
     "import numpyro.distributions as dist\n",
     "from numpyro.infer import MCMC, NUTS\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.1\")\n",
     "\n",
     "# NB: replace cpu by gpu to run this notebook on gpu\n",
     "numpyro.set_platform(\"cpu\")"
diff --git a/notebooks/source/bayesian_hierarchical_linear_regression.ipynb b/notebooks/source/bayesian_hierarchical_linear_regression.ipynb
index 8d29d9158..36d0f743f 100644
--- a/notebooks/source/bayesian_hierarchical_linear_regression.ipynb
+++ b/notebooks/source/bayesian_hierarchical_linear_regression.ipynb
@@ -246,7 +246,7 @@
     "import numpyro.distributions as dist\n",
     "from numpyro.infer import MCMC, NUTS, Predictive\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_hierarchical_stacking.ipynb b/notebooks/source/bayesian_hierarchical_stacking.ipynb
index e03174ef7..dff9d67ad 100644
--- a/notebooks/source/bayesian_hierarchical_stacking.ipynb
+++ b/notebooks/source/bayesian_hierarchical_stacking.ipynb
@@ -97,7 +97,7 @@
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
     "numpyro.set_host_device_count(4)\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_imputation.ipynb b/notebooks/source/bayesian_imputation.ipynb
index ab223bdf1..b1b0381da 100644
--- a/notebooks/source/bayesian_imputation.ipynb
+++ b/notebooks/source/bayesian_imputation.ipynb
@@ -52,7 +52,7 @@
     "if \"NUMPYRO_SPHINXBUILD\" in os.environ:\n",
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_regression.ipynb b/notebooks/source/bayesian_regression.ipynb
index f1a8bcf20..1c7a8ae40 100644
--- a/notebooks/source/bayesian_regression.ipynb
+++ b/notebooks/source/bayesian_regression.ipynb
@@ -73,7 +73,7 @@
     "if \"NUMPYRO_SPHINXBUILD\" in os.environ:\n",
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/censoring.ipynb b/notebooks/source/censoring.ipynb
index a1c2d6dff..204a02c1f 100644
--- a/notebooks/source/censoring.ipynb
+++ b/notebooks/source/censoring.ipynb
@@ -60,7 +60,7 @@
     "\n",
     "rng_key = random.PRNGKey(seed=0)\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.1\")\n",
     "\n",
     "%load_ext autoreload\n",
     "%autoreload 2\n",
diff --git a/notebooks/source/gmm.ipynb b/notebooks/source/gmm.ipynb
index b64e65167..f45a07b1f 100644
--- a/notebooks/source/gmm.ipynb
+++ b/notebooks/source/gmm.ipynb
@@ -54,7 +54,7 @@
     "%matplotlib inline\n",
     "\n",
     "smoke_test = \"CI\" in os.environ\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/hsgp_example.ipynb b/notebooks/source/hsgp_example.ipynb
index 0bfcafc23..d5bef357f 100644
--- a/notebooks/source/hsgp_example.ipynb
+++ b/notebooks/source/hsgp_example.ipynb
@@ -62,7 +62,7 @@
     "\n",
     "rng_key = random.PRNGKey(seed=42)\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.1\")\n",
     "\n",
     "%load_ext autoreload\n",
     "%autoreload 2\n",
diff --git a/notebooks/source/logistic_regression.ipynb b/notebooks/source/logistic_regression.ipynb
index 0432ca010..917004657 100644
--- a/notebooks/source/logistic_regression.ipynb
+++ b/notebooks/source/logistic_regression.ipynb
@@ -41,7 +41,7 @@
     "from numpyro.examples.datasets import COVTYPE, load_dataset\n",
     "from numpyro.infer import HMC, MCMC, NUTS\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.1\")\n",
     "\n",
     "# NB: replace gpu by cpu to run this notebook in cpu\n",
     "numpyro.set_platform(\"gpu\")"
diff --git a/notebooks/source/model_rendering.ipynb b/notebooks/source/model_rendering.ipynb
index 359d6886a..d06abaf33 100644
--- a/notebooks/source/model_rendering.ipynb
+++ b/notebooks/source/model_rendering.ipynb
@@ -38,7 +38,7 @@
     "import numpyro.distributions as dist\n",
     "import numpyro.distributions.constraints as constraints\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/ordinal_regression.ipynb b/notebooks/source/ordinal_regression.ipynb
index fcb610095..174de9f9c 100644
--- a/notebooks/source/ordinal_regression.ipynb
+++ b/notebooks/source/ordinal_regression.ipynb
@@ -54,7 +54,7 @@
     "from numpyro.infer import MCMC, NUTS\n",
     "from numpyro.infer.reparam import TransformReparam\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/notebooks/source/time_series_forecasting.ipynb b/notebooks/source/time_series_forecasting.ipynb
index 40f467ee0..f535cecd5 100644
--- a/notebooks/source/time_series_forecasting.ipynb
+++ b/notebooks/source/time_series_forecasting.ipynb
@@ -48,7 +48,7 @@
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
     "numpyro.set_host_device_count(4)\n",
-    "assert numpyro.__version__.startswith(\"0.15.0\")"
+    "assert numpyro.__version__.startswith(\"0.15.1\")"
    ]
   },
   {
diff --git a/numpyro/version.py b/numpyro/version.py
index c44b61ed5..4618c3c4f 100644
--- a/numpyro/version.py
+++ b/numpyro/version.py
@@ -1,4 +1,4 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "0.15.0"
+__version__ = "0.15.1"


Test patch:
diff --git a/examples/proportion_test.py b/examples/proportion_test.py
index bdbe1330f..ccf5a6c61 100644
--- a/examples/proportion_test.py
+++ b/examples/proportion_test.py
@@ -158,7 +158,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.15.0")
+    assert numpyro.__version__.startswith("0.15.1")
     parser = argparse.ArgumentParser(description="Testing whether  ")
     parser.add_argument("-n", "--num-samples", nargs="?", default=500, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1500, type=int)