Bump to 0.9.0 (#1310)

* Add loose strategy for MCMC

* merge svi and mcmc plate warning strategies

* fix failing tests

* validate model accross ELBOs

* update vae example

* fix typos

* Bump to version 0.9.0

* Fix failing tests

* Fix warnings in tests/examples

* relax funsor requirement

* Move optax_to_numpyro to optim

* skip prodlda test on CI

Author: fehiepsi

.github/workflows/ci.yml

@@ -26,10 +26,10 @@ jobs:
       run: |
         sudo apt install -y pandoc gsfonts
         python -m pip install --upgrade pip
-        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install jaxlib
         pip install jax
         pip install .[doc,test]
+        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install -r docs/requirements.txt
         pip freeze
     - name: Lint with flake8
@@ -64,10 +64,10 @@ jobs:
         python -m pip install --upgrade pip
         # Keep track of pyro-api master branch
         pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip
-        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install jaxlib
         pip install jax
         pip install .[dev,test]
+        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip freeze
     - name: Test with pytest
       run: |
@@ -93,10 +93,10 @@ jobs:
         python -m pip install --upgrade pip
         # Keep track of pyro-api master branch
         pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip
-        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install jaxlib
         pip install jax
         pip install .[dev,test]
+        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip freeze
     - name: Test with pytest
       run: |
@@ -129,10 +129,10 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install jaxlib
         pip install jax
         pip install .[dev,examples,test]
+        pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip freeze
     - name: Test with pytest
       run: |

docs/source/optimizers.rst

@@ -68,4 +68,4 @@ SM3
 
 Optax support
 -------------
-.. autofunction:: numpyro.contrib.optim.optax_to_numpyro
+.. autofunction:: numpyro.optim.optax_to_numpyro

examples/annotation.py

@@ -351,7 +351,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.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)

examples/ar2.py

@@ -116,7 +116,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.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)

examples/baseball.py

@@ -210,7 +210,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.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)

examples/bnn.py

@@ -160,7 +160,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.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)

examples/capture_recapture.py

@@ -72,7 +72,11 @@ def transition_fn(carry, y):
             with handlers.mask(mask=first_capture_mask):
                 mu_z_t = first_capture_mask * phi * z + (1 - first_capture_mask)
                 # NumPyro exactly sums out the discrete states z_t.
-                z = numpyro.sample("z", dist.Bernoulli(dist.util.clamp_probs(mu_z_t)))
+                z = numpyro.sample(
+                    "z",
+                    dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
+                    infer={"enumerate": "parallel"},
+                )
                 mu_y_t = rho * z
                 numpyro.sample(
                     "y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
@@ -112,7 +116,11 @@ def transition_fn(carry, y):
             with handlers.mask(mask=first_capture_mask):
                 mu_z_t = first_capture_mask * phi_t * z + (1 - first_capture_mask)
                 # NumPyro exactly sums out the discrete states z_t.
-                z = numpyro.sample("z", dist.Bernoulli(dist.util.clamp_probs(mu_z_t)))
+                z = numpyro.sample(
+                    "z",
+                    dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
+                    infer={"enumerate": "parallel"},
+                )
                 mu_y_t = rho * z
                 numpyro.sample(
                     "y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
@@ -160,7 +168,11 @@ def transition_fn(carry, y):
             with handlers.mask(mask=first_capture_mask):
                 mu_z_t = first_capture_mask * phi_t * z + (1 - first_capture_mask)
                 # NumPyro exactly sums out the discrete states z_t.
-                z = numpyro.sample("z", dist.Bernoulli(dist.util.clamp_probs(mu_z_t)))
+                z = numpyro.sample(
+                    "z",
+                    dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
+                    infer={"enumerate": "parallel"},
+                )
                 mu_y_t = rho * z
                 numpyro.sample(
                     "y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
@@ -202,7 +214,11 @@ def transition_fn(carry, y):
             with handlers.mask(mask=first_capture_mask):
                 mu_z_t = first_capture_mask * phi * z + (1 - first_capture_mask)
                 # NumPyro exactly sums out the discrete states z_t.
-                z = numpyro.sample("z", dist.Bernoulli(dist.util.clamp_probs(mu_z_t)))
+                z = numpyro.sample(
+                    "z",
+                    dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
+                    infer={"enumerate": "parallel"},
+                )
                 mu_y_t = rho * z
                 numpyro.sample(
                     "y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y
@@ -249,7 +265,11 @@ def transition_fn(carry, y):
             with handlers.mask(mask=first_capture_mask):
                 mu_z_t = first_capture_mask * phi_t * z + (1 - first_capture_mask)
                 # NumPyro exactly sums out the discrete states z_t.
-                z = numpyro.sample("z", dist.Bernoulli(dist.util.clamp_probs(mu_z_t)))
+                z = numpyro.sample(
+                    "z",
+                    dist.Bernoulli(dist.util.clamp_probs(mu_z_t)),
+                    infer={"enumerate": "parallel"},
+                )
                 mu_y_t = rho * z
                 numpyro.sample(
                     "y", dist.Bernoulli(dist.util.clamp_probs(mu_y_t)), obs=y

examples/covtype.py

@@ -206,7 +206,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.0")
     parser = argparse.ArgumentParser(description="parse args")
     parser.add_argument(
         "-n", "--num-samples", default=1000, type=int, help="number of samples"

examples/funnel.py

@@ -139,7 +139,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.0")
     parser = argparse.ArgumentParser(
         description="Non-centered reparameterization example"
     )

examples/gaussian_shells.py

@@ -81,7 +81,7 @@ def run_inference(args, data):
             num_warmup=args.num_warmup,
             num_samples=args.num_samples,
         )
-    mcmc.run(random.PRNGKey(2), **data)
+    mcmc.run(random.PRNGKey(2), **data, enum=args.enum)
     mcmc.print_summary()
     mcmc_samples = mcmc.get_samples()
 
@@ -123,7 +123,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.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)