Bump to 0.13.0 (#1629)

* bump jax version to 0.4.14 * Bump to 0.13.0 * drop 3.8 in setup * update typing for python 3.9 * more typing from collections.abc
pyro-ppl · Aug 21, 2023 · d8e1156 · d8e1156
1 parent cc1ad02
commit d8e1156
Show file tree

Hide file tree

Showing 53 changed files with 81 additions and 159 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8]
+        python-version: [3.9]
 
     steps:
     - uses: actions/checkout@v2
@@ -51,7 +51,7 @@ jobs:
     needs: lint
     strategy:
       matrix:
-        python-version: [3.8]
+        python-version: [3.9]
 
     steps:
     - uses: actions/checkout@v2
@@ -81,7 +81,7 @@ jobs:
     needs: lint
     strategy:
       matrix:
-        python-version: [3.8]
+        python-version: [3.9]
 
     steps:
     - uses: actions/checkout@v2
@@ -120,7 +120,7 @@ jobs:
     needs: lint
     strategy:
       matrix:
-        python-version: [3.8]
+        python-version: [3.9]
 
     steps:
     - uses: actions/checkout@v2

diff --git a/examples/annotation.py b/examples/annotation.py
@@ -320,7 +320,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -114,7 +114,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -210,7 +210,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -160,7 +160,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -206,7 +206,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -139,7 +139,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     parser = argparse.ArgumentParser(
         description="Non-centered reparameterization example"
     )

diff --git 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.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -170,7 +170,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -263,7 +263,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -180,7 +180,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -162,7 +162,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -58,7 +58,7 @@ def body_fn(i, val):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -220,7 +220,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
 
     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
@@ -197,7 +197,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -117,7 +117,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -314,7 +314,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     parser = argparse.ArgumentParser(
         description="Probabilistic topic modelling with Flax and Haiku"
     )

diff --git a/examples/proportion_test.py b/examples/proportion_test.py
@@ -19,7 +19,6 @@
 
 import argparse
 import os
-from typing import Tuple
 
 from jax import random
 import jax.numpy as jnp
@@ -31,7 +30,7 @@
 from numpyro.infer import MCMC, NUTS
 
 
-def make_dataset(rng_key) -> Tuple[jnp.ndarray, jnp.ndarray]:
+def make_dataset(rng_key) -> tuple[jnp.ndarray, jnp.ndarray]:
     """
     Make simulated dataset where potential customers who get a
     sales calls have ~2% higher chance of making another purchase.
@@ -160,7 +159,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -384,7 +384,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -122,7 +122,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -292,7 +292,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -126,7 +126,7 @@ def get_true_pred_CPDs(CPD, posterior_param):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -151,7 +151,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     parser = argparse.ArgumentParser(
         description="UCBadmit gender discrimination using HMC"
     )

diff --git 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.12.1")
+    assert numpyro.__version__.startswith("0.13.0")
     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
@@ -50,7 +50,7 @@
     "\n",
     "from numpyro.infer import MCMC, NUTS\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.12.1\")\n",
+    "assert numpyro.__version__.startswith(\"0.13.0\")\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
@@ -245,7 +245,7 @@
     "import numpyro.distributions as dist\n",
     "from jax import random\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git 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.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git 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.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git 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.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ],
    "execution_count": 2,
    "outputs": []

diff --git 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.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git 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.12.1\")\n",
+    "assert numpyro.__version__.startswith(\"0.13.0\")\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
@@ -37,7 +37,7 @@
     "import numpyro.distributions as dist\n",
     "import numpyro.distributions.constraints as constraints\n",
     "\n",
-    "assert numpyro.__version__.startswith(\"0.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git 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.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git 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.12.1\")"
+    "assert numpyro.__version__.startswith(\"0.13.0\")"
    ]
   },
   {

diff --git a/numpyro/contrib/einstein/mixture_guide_predictive.py b/numpyro/contrib/einstein/mixture_guide_predictive.py
@@ -1,8 +1,9 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from collections.abc import Callable, Sequence
 from functools import partial
-from typing import Callable, Dict, Optional, Sequence
+from typing import Optional
 
 from jax import numpy as jnp, random, tree_map, vmap
 from jax.tree_util import tree_flatten
@@ -24,7 +25,7 @@ class MixtureGuidePredictive:
 
     :param Callable model: Python callable containing Pyro primitives.
     :param Callable guide: Python callable containing Pyro primitives to get posterior samples of sites.
-    :param Dict params:  Dictionary of values for param sites of model/guide
+    :param dict params:  Dictionary of values for param sites of model/guide
     :param Sequence guide_sites: Names of sites that contribute to the Stein mixture.
     :param Optional[int] num_samples:
     :param Optional[Sequence[str]] return_sites: Sites to return. By default, only sample sites not present
@@ -37,7 +38,7 @@ def __init__(
         self,
         model: Callable,
         guide: Callable,
-        params: Dict,
+        params: dict,
         guide_sites: Sequence,
         num_samples: Optional[int] = None,
         return_sites: Optional[Sequence[str]] = None,