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

.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)

examples/gp.py

@@ -170,7 +170,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="Gaussian Process example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)

examples/hmm.py

@@ -263,7 +263,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="Semi-supervised Hidden Markov Model")
     parser.add_argument("--num-categories", default=3, type=int)
     parser.add_argument("--num-words", default=10, type=int)

examples/hmm_enum.py

@@ -94,7 +94,11 @@ def transition_fn(carry, y):
         x_prev, t = carry
         with numpyro.plate("sequences", num_sequences, dim=-2):
             with mask(mask=(t < lengths)[..., None]):
-                x = numpyro.sample("x", dist.Categorical(probs_x[x_prev]))
+                x = numpyro.sample(
+                    "x",
+                    dist.Categorical(probs_x[x_prev]),
+                    infer={"enumerate": "parallel"},
+                )
                 with numpyro.plate("tones", data_dim, dim=-1):
                     numpyro.sample("y", dist.Bernoulli(probs_y[x.squeeze(-1)]), obs=y)
         return (x, t + 1), None
@@ -127,7 +131,11 @@ def transition_fn(carry, y):
         x_prev, y_prev, t = carry
         with numpyro.plate("sequences", num_sequences, dim=-2):
             with mask(mask=(t < lengths)[..., None]):
-                x = numpyro.sample("x", dist.Categorical(probs_x[x_prev]))
+                x = numpyro.sample(
+                    "x",
+                    dist.Categorical(probs_x[x_prev]),
+                    infer={"enumerate": "parallel"},
+                )
                 # Note the broadcasting tricks here: to index probs_y on tensors x and y,
                 # we also need a final tensor for the tones dimension. This is conveniently
                 # provided by the plate associated with that dimension.
@@ -175,8 +183,16 @@ def transition_fn(carry, y):
         w_prev, x_prev, t = carry
         with numpyro.plate("sequences", num_sequences, dim=-2):
             with mask(mask=(t < lengths)[..., None]):
-                w = numpyro.sample("w", dist.Categorical(probs_w[w_prev]))
-                x = numpyro.sample("x", dist.Categorical(probs_x[x_prev]))
+                w = numpyro.sample(
+                    "w",
+                    dist.Categorical(probs_w[w_prev]),
+                    infer={"enumerate": "parallel"},
+                )
+                x = numpyro.sample(
+                    "x",
+                    dist.Categorical(probs_x[x_prev]),
+                    infer={"enumerate": "parallel"},
+                )
                 # Note the broadcasting tricks here: to index probs_y on tensors x and y,
                 # we also need a final tensor for the tones dimension. This is conveniently
                 # provided by the plate associated with that dimension.
@@ -224,8 +240,16 @@ def transition_fn(carry, y):
         w_prev, x_prev, t = carry
         with numpyro.plate("sequences", num_sequences, dim=-2):
             with mask(mask=(t < lengths)[..., None]):
-                w = numpyro.sample("w", dist.Categorical(probs_w[w_prev]))
-                x = numpyro.sample("x", dist.Categorical(Vindex(probs_x)[w, x_prev]))
+                w = numpyro.sample(
+                    "w",
+                    dist.Categorical(probs_w[w_prev]),
+                    infer={"enumerate": "parallel"},
+                )
+                x = numpyro.sample(
+                    "x",
+                    dist.Categorical(Vindex(probs_x)[w, x_prev]),
+                    infer={"enumerate": "parallel"},
+                )
                 with numpyro.plate("tones", data_dim, dim=-1) as tones:
                     numpyro.sample("y", dist.Bernoulli(probs_y[w, x, tones]), obs=y)
         return (w, x, t + 1), None
@@ -275,7 +299,7 @@ def transition_fn(carry, y):
             with mask(mask=(t < lengths)[..., None]):
                 probs_x_t = Vindex(probs_x)[x_prev, x_curr]
                 x_prev, x_curr = x_curr, numpyro.sample(
-                    "x", dist.Categorical(probs_x_t)
+                    "x", dist.Categorical(probs_x_t), infer={"enumerate": "parallel"}
                 )
                 with numpyro.plate("tones", data_dim, dim=-1):
                     probs_y_t = probs_y[x_curr.squeeze(-1)]

examples/holt_winters.py

@@ -180,7 +180,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="Holt-Winters")
     parser.add_argument("--T", nargs="?", default=6, type=int)
     parser.add_argument("--future", nargs="?", default=1, type=int)

examples/horseshoe_regression.py

@@ -162,7 +162,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="Horseshoe regression example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=2000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)

examples/hsgp.py

@@ -108,7 +108,7 @@ def get_thanksgiving_days(dates):
 
 def get_floating_days_indicators(dates):
     def encode(x):
-        return jnp.array(x.values, dtype=jnp.int64)
+        return jnp.array(x.values, dtype=jnp.result_type(int))
 
     return {
         "labour_days_indicator": encode(get_labour_days(dates)),

examples/minipyro.py

@@ -58,7 +58,7 @@ def body_fn(i, val):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.8.0")
+    assert numpyro.__version__.startswith("0.9.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)

examples/neutra.py

@@ -197,7 +197,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="NeuTra HMC")
     parser.add_argument("-n", "--num-samples", nargs="?", default=4000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)

examples/ode.py

@@ -117,7 +117,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="Predator-Prey Model")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)

examples/prodlda.py

@@ -314,7 +314,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="Probabilistic topic modelling with Flax and Haiku"
     )

examples/proportion_test.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="Testing whether  ")
     parser.add_argument("-n", "--num-samples", nargs="?", default=500, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1500, type=int)

examples/sparse_regression.py

@@ -384,7 +384,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="Gaussian Process example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=500, type=int)

examples/stochastic_volatility.py

@@ -122,7 +122,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="Stochastic Volatility Model")
     parser.add_argument("-n", "--num-samples", nargs="?", default=600, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=600, type=int)

examples/thompson_sampling.py

@@ -294,7 +294,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="Thompson sampling example")
     parser.add_argument(
         "--num-random", nargs="?", default=2, type=int, help="number of random draws"

examples/ucbadmit.py

@@ -151,7 +151,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="UCBadmit gender discrimination using HMC"
     )

examples/vae.py

@@ -160,7 +160,7 @@ def reconstruct_img(epoch, rng_key):
 
 
 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-epochs", default=15, type=int, help="number of training epochs"