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

Issue:
Simplify the provenance logic to prepare for the removal of jax named shape
See https://github.com/google/jax/pull/21069 as a context. The current tests fail with jax's dev branch.

Current tests at `test/test_model_rendering` and `test/ops/test_provenance` pass locally.


Discussion / hints:
Yes, this just simplifies the logic. Previously we used named shape because we need some functionality to flatten unflatten the provenance. Turns out that those jax tree utilities also work for frozenset leaves.

Fix patch:
diff --git a/examples/annotation.py b/examples/annotation.py
index a28ab48be..119409d30 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 6ebb5949e..a73844c56 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 47615c37a..b0a08ea25 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 022e88d9e..4376d88c6 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 424eac7f7..a5a972323 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 bb28e7dfa..988ff2386 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     parser = argparse.ArgumentParser(
         description="Non-centered reparameterization example"
     )
diff --git a/examples/gaussian_shells.py b/examples/gaussian_shells.py
index 9c2184eb4..f0af3c095 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 dfa22f624..4c77b4cf7 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 3d17b9366..8d1290475 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 fb97e62bd..946aec23e 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 ee1624cc2..081983eee 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 8d7094de3..c99cfb26b 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 0ee18dc06..406ece20e 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
 
     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 c47d8f94e..8ac495767 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 3845dcd9c..6db17df49 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 0a30d202d..78d272061 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     parser = argparse.ArgumentParser(
         description="Probabilistic topic modelling with Flax and Haiku"
     )
diff --git a/examples/sparse_regression.py b/examples/sparse_regression.py
index f0c9781d4..48ea44d33 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 a4ce6d7fd..c1fafd232 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 e10205b72..906d71376 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 f64bdb98b..87a4df6ce 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 4147067c6..c4afe784f 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     parser = argparse.ArgumentParser(
         description="UCBadmit gender discrimination using HMC"
     )
diff --git a/examples/vae.py b/examples/vae.py
index 5543c30d0..446e5a19f 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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 68c2f0c56..ddffcda44 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.1\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.2\")\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 36d0f743f..cfa800d0e 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_hierarchical_stacking.ipynb b/notebooks/source/bayesian_hierarchical_stacking.ipynb
index dff9d67ad..e11e0a3dc 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_imputation.ipynb b/notebooks/source/bayesian_imputation.ipynb
index b1b0381da..f4746341a 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_regression.ipynb b/notebooks/source/bayesian_regression.ipynb
index 1c7a8ae40..6adda4aa9 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/censoring.ipynb b/notebooks/source/censoring.ipynb
index 204a02c1f..2deed7378 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.1\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.2\")\n",
     "\n",
     "%load_ext autoreload\n",
     "%autoreload 2\n",
diff --git a/notebooks/source/gmm.ipynb b/notebooks/source/gmm.ipynb
index f45a07b1f..129bffede 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/hsgp_example.ipynb b/notebooks/source/hsgp_example.ipynb
index d5bef357f..881bd0743 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.1\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.2\")\n",
     "\n",
     "%load_ext autoreload\n",
     "%autoreload 2\n",
diff --git a/notebooks/source/logistic_regression.ipynb b/notebooks/source/logistic_regression.ipynb
index 917004657..6ea78044c 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.1\")\n",
+    "assert numpyro.__version__.startswith(\"0.15.2\")\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 d06abaf33..377dfdaa7 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/ordinal_regression.ipynb b/notebooks/source/ordinal_regression.ipynb
index 174de9f9c..ad7f8340a 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/notebooks/source/time_series_forecasting.ipynb b/notebooks/source/time_series_forecasting.ipynb
index f535cecd5..3744ebe0b 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.1\")"
+    "assert numpyro.__version__.startswith(\"0.15.2\")"
    ]
   },
   {
diff --git a/numpyro/version.py b/numpyro/version.py
index 4618c3c4f..5ff26e807 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.1"
+__version__ = "0.15.2"


Test patch:
diff --git a/examples/proportion_test.py b/examples/proportion_test.py
index ccf5a6c61..477c5695b 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.1")
+    assert numpyro.__version__.startswith("0.15.2")
     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)