From 3c199d2eedd15f24fcb0975670e91b84d8405713 Mon Sep 17 00:00:00 2001
From: Du Phan <phandu@google.com>
Date: Mon, 29 Aug 2022 14:23:45 -0400
Subject: [PATCH] Bump to version 0.10.1 (#1474)

---
 examples/annotation.py                                     | 2 +-
 examples/ar2.py                                            | 2 +-
 examples/baseball.py                                       | 2 +-
 examples/bnn.py                                            | 2 +-
 examples/covtype.py                                        | 2 +-
 examples/funnel.py                                         | 2 +-
 examples/gaussian_shells.py                                | 2 +-
 examples/gp.py                                             | 2 +-
 examples/hmm.py                                            | 2 +-
 examples/holt_winters.py                                   | 2 +-
 examples/horseshoe_regression.py                           | 2 +-
 examples/minipyro.py                                       | 2 +-
 examples/mortality.py                                      | 2 +-
 examples/neutra.py                                         | 2 +-
 examples/ode.py                                            | 2 +-
 examples/prodlda.py                                        | 2 +-
 examples/proportion_test.py                                | 2 +-
 examples/sparse_regression.py                              | 2 +-
 examples/stochastic_volatility.py                          | 2 +-
 examples/thompson_sampling.py                              | 2 +-
 examples/ucbadmit.py                                       | 2 +-
 examples/vae.py                                            | 2 +-
 notebooks/source/bad_posterior_geometry.ipynb              | 2 +-
 .../source/bayesian_hierarchical_linear_regression.ipynb   | 2 +-
 notebooks/source/bayesian_hierarchical_stacking.ipynb      | 2 +-
 notebooks/source/bayesian_imputation.ipynb                 | 2 +-
 notebooks/source/bayesian_regression.ipynb                 | 2 +-
 notebooks/source/logistic_regression.ipynb                 | 2 +-
 notebooks/source/model_rendering.ipynb                     | 2 +-
 notebooks/source/ordinal_regression.ipynb                  | 2 +-
 notebooks/source/time_series_forecasting.ipynb             | 2 +-
 numpyro/infer/svi.py                                       | 2 +-
 numpyro/optim.py                                           | 2 +-
 numpyro/version.py                                         | 2 +-
 test/infer/test_svi.py                                     | 2 +-
 test/test_handlers.py                                      | 7 +++++++
 36 files changed, 42 insertions(+), 35 deletions(-)

diff --git a/examples/annotation.py b/examples/annotation.py
index aa7f6af4e..30fe16e1e 100644
--- a/examples/annotation.py
+++ b/examples/annotation.py
@@ -320,7 +320,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 3b8602ab7..a98bd07c9 100644
--- a/examples/ar2.py
+++ b/examples/ar2.py
@@ -114,7 +114,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 8804d54d8..509bc5450 100644
--- a/examples/baseball.py
+++ b/examples/baseball.py
@@ -210,7 +210,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 92ffa2d15..9e8c85390 100644
--- a/examples/bnn.py
+++ b/examples/bnn.py
@@ -160,7 +160,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 b8819a513..44059d68e 100644
--- a/examples/covtype.py
+++ b/examples/covtype.py
@@ -206,7 +206,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 9fbe0c0c2..b5fcfeb49 100644
--- a/examples/funnel.py
+++ b/examples/funnel.py
@@ -139,7 +139,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.1")
     parser = argparse.ArgumentParser(
         description="Non-centered reparameterization example"
     )
diff --git a/examples/gaussian_shells.py b/examples/gaussian_shells.py
index 78439f16f..84dde27d5 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 741e783ad..4c3cddb31 100644
--- a/examples/gp.py
+++ b/examples/gp.py
@@ -170,7 +170,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 71adc3b5b..c336eb41f 100644
--- a/examples/hmm.py
+++ b/examples/hmm.py
@@ -263,7 +263,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 bfed87a6b..e80ebe20c 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 f1799f7fd..95daf9714 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 445691980..2d3015d42 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 7e57e5d47..3ce68df4c 100644
--- a/examples/mortality.py
+++ b/examples/mortality.py
@@ -220,7 +220,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 4e7f7cd93..9709e436f 100644
--- a/examples/neutra.py
+++ b/examples/neutra.py
@@ -197,7 +197,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 d06b52a48..1e00a091f 100644
--- a/examples/ode.py
+++ b/examples/ode.py
@@ -117,7 +117,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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 5eba7631e..9f3471f19 100644
--- a/examples/prodlda.py
+++ b/examples/prodlda.py
@@ -314,7 +314,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.1")
     parser = argparse.ArgumentParser(
         description="Probabilistic topic modelling with Flax and Haiku"
     )
diff --git a/examples/proportion_test.py b/examples/proportion_test.py
index cc4276d23..5f2fc9429 100644
--- a/examples/proportion_test.py
+++ b/examples/proportion_test.py
@@ -160,7 +160,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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)
diff --git a/examples/sparse_regression.py b/examples/sparse_regression.py
index a52a6938b..12728a968 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 e504b1ae3..255d15c1c 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 4b5f74280..55638eb24 100644
--- a/examples/thompson_sampling.py
+++ b/examples/thompson_sampling.py
@@ -294,7 +294,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.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/ucbadmit.py b/examples/ucbadmit.py
index 960ad4755..0042f9e44 100644
--- a/examples/ucbadmit.py
+++ b/examples/ucbadmit.py
@@ -151,7 +151,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.10.0")
+    assert numpyro.__version__.startswith("0.10.1")
     parser = argparse.ArgumentParser(
         description="UCBadmit gender discrimination using HMC"
     )
diff --git a/examples/vae.py b/examples/vae.py
index 6386b8356..308515113 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.10.0")
+    assert numpyro.__version__.startswith("0.10.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 fe303e730..05eab1cb4 100644
--- a/notebooks/source/bad_posterior_geometry.ipynb
+++ b/notebooks/source/bad_posterior_geometry.ipynb
@@ -50,7 +50,7 @@
     "\n",
     "from numpyro.infer import MCMC, NUTS\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")\n",
+    "assert numpyro.__version__.startswith(\"0.10.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 6135bdd88..905b16a68 100644
--- a/notebooks/source/bayesian_hierarchical_linear_regression.ipynb
+++ b/notebooks/source/bayesian_hierarchical_linear_regression.ipynb
@@ -244,7 +244,7 @@
     "import numpyro.distributions as dist\n",
     "from jax import random\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_hierarchical_stacking.ipynb b/notebooks/source/bayesian_hierarchical_stacking.ipynb
index 4382665f4..b819beaec 100644
--- a/notebooks/source/bayesian_hierarchical_stacking.ipynb
+++ b/notebooks/source/bayesian_hierarchical_stacking.ipynb
@@ -96,7 +96,7 @@
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
     "numpyro.set_host_device_count(4)\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_imputation.ipynb b/notebooks/source/bayesian_imputation.ipynb
index 8bef3cec4..e979a6286 100644
--- a/notebooks/source/bayesian_imputation.ipynb
+++ b/notebooks/source/bayesian_imputation.ipynb
@@ -55,7 +55,7 @@
     "if \"NUMPYRO_SPHINXBUILD\" in os.environ:\n",
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ]
   },
   {
diff --git a/notebooks/source/bayesian_regression.ipynb b/notebooks/source/bayesian_regression.ipynb
index b760ef4de..b45454528 100644
--- a/notebooks/source/bayesian_regression.ipynb
+++ b/notebooks/source/bayesian_regression.ipynb
@@ -95,7 +95,7 @@
     "if \"NUMPYRO_SPHINXBUILD\" in os.environ:\n",
     "    set_matplotlib_formats(\"svg\")\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ],
    "execution_count": 2,
    "outputs": []
diff --git a/notebooks/source/logistic_regression.ipynb b/notebooks/source/logistic_regression.ipynb
index e9b856444..2daf0d837 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.10.0\")\n",
+    "assert numpyro.__version__.startswith(\"0.10.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 e12abaf39..974c07bf1 100644
--- a/notebooks/source/model_rendering.ipynb
+++ b/notebooks/source/model_rendering.ipynb
@@ -37,7 +37,7 @@
     "import numpyro.distributions as dist\n",
     "import numpyro.distributions.constraints as constraints\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ]
   },
   {
diff --git a/notebooks/source/ordinal_regression.ipynb b/notebooks/source/ordinal_regression.ipynb
index 4e5963a4f..f7deb8621 100644
--- a/notebooks/source/ordinal_regression.ipynb
+++ b/notebooks/source/ordinal_regression.ipynb
@@ -53,7 +53,7 @@
     "import pandas as pd\n",
     "import seaborn as sns\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ]
   },
   {
diff --git a/notebooks/source/time_series_forecasting.ipynb b/notebooks/source/time_series_forecasting.ipynb
index 166439ec2..24bd681cb 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.10.0\")"
+    "assert numpyro.__version__.startswith(\"0.10.1\")"
    ]
   },
   {
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
index 6702d9079..017869a6b 100644
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -13,7 +13,7 @@
 if _versiontuple(jax.__version__) >= (0, 2, 25):
     from jax.example_libraries import optimizers
 else:
-    from jax.experimental import optimizers
+    from jax.experimental import optimizers  # pytype: disable=import-error
 
 from jax import jit, lax, random
 import jax.numpy as jnp
diff --git a/numpyro/optim.py b/numpyro/optim.py
index b7bd54efd..00d08906f 100644
--- a/numpyro/optim.py
+++ b/numpyro/optim.py
@@ -18,7 +18,7 @@
 if _versiontuple(jax.__version__) >= (0, 2, 25):
     from jax.example_libraries import optimizers
 else:
-    from jax.experimental import optimizers
+    from jax.experimental import optimizers  # pytype: disable=import-error
 
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
diff --git a/numpyro/version.py b/numpyro/version.py
index 4f527df18..8a73a33bb 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.10.0"
+__version__ = "0.10.1"
diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py
index d994e5ce2..16d227722 100644
--- a/test/infer/test_svi.py
+++ b/test/infer/test_svi.py
@@ -17,7 +17,7 @@
 if _versiontuple(jax.__version__) >= (0, 2, 25):
     from jax.example_libraries import optimizers
 else:
-    from jax.experimental import optimizers
+    from jax.experimental import optimizers  # pytype: disable=import-error
 
 import numpyro
 from numpyro import optim
diff --git a/test/test_handlers.py b/test/test_handlers.py
index f40567597..514ecdeb0 100644
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -10,6 +10,11 @@
 import jax.numpy as jnp
 from jax.tree_util import tree_map
 
+try:
+    import funsor
+except ImportError:
+    funsor = None
+
 import numpyro
 from numpyro import handlers
 import numpyro.distributions as dist
@@ -658,6 +663,7 @@ def model():
             model()
 
 
+@pytest.mark.skipif(funsor is None, reason="require funsor installation")
 def test_collapse_beta_binomial():
     total_count = 10
     data = 3.0
@@ -696,6 +702,7 @@ def model2():
     assert_allclose(params1["c0"], params2["c0"])
 
 
+@pytest.mark.skipif(funsor is None, reason="require funsor installation")
 def test_collapse_beta_bernoulli():
     data = 0.0