explicit num_warmup, num_samples arguments for MCMC (#1040)

fehiepsi · web-flow · commit d735523972ae · 2021-05-26T12:21:55.000-05:00
* request for explicit num_warmup, num_samples

* force num_warmup num_samples in all tests

* ping nbsphinx version

* fix remaining bugs

* run make format

* revert docs requirement

* fix wrong changes
diff --git a/examples/annotation.py b/examples/annotation.py
@@ -304,8 +304,8 @@ def main(args):
 
     mcmc = MCMC(
         NUTS(model),
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/baseball.py b/examples/baseball.py
@@ -145,8 +145,8 @@ def run_inference(model, at_bats, hits, rng_key, args):
         kernel = SA(model)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False
         if ("NUMPYRO_SPHINXBUILD" in os.environ or args.disable_progbar)
diff --git a/examples/bnn.py b/examples/bnn.py
@@ -76,8 +76,8 @@ def run_inference(model, args, rng_key, X, Y, D_H):
     kernel = NUTS(model)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/capture_recapture.py b/examples/capture_recapture.py
@@ -288,8 +288,8 @@ def run_inference(model, capture_history, sex, rng_key, args):
         kernel = HMC(model)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/covtype.py b/examples/covtype.py
@@ -192,7 +192,7 @@ def benchmark_hmc(args, features, labels):
         )
     else:
         raise ValueError("Invalid algorithm, either 'HMC', 'NUTS', or 'HMCECS'.")
-    mcmc = MCMC(kernel, args.num_warmup, args.num_samples)
+    mcmc = MCMC(kernel, num_warmup=args.num_warmup, num_samples=args.num_samples)
     mcmc.run(rng_key, features, labels, subsample_size, extra_fields=("accept_prob",))
     print("Mean accept prob:", jnp.mean(mcmc.get_extra_fields()["accept_prob"]))
     mcmc.print_summary(exclude_deterministic=False)
diff --git a/examples/funnel.py b/examples/funnel.py
@@ -53,8 +53,8 @@ def run_inference(model, args, rng_key):
     kernel = NUTS(model)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/gp.py b/examples/gp.py
@@ -85,8 +85,8 @@ def run_inference(model, args, rng_key, X, Y):
     kernel = NUTS(model, init_strategy=init_strategy)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         thinning=args.thinning,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
diff --git a/examples/hmm.py b/examples/hmm.py
@@ -221,8 +221,8 @@ def main(args):
     kernel = NUTS(semi_supervised_hmm)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/hmm_enum.py b/examples/hmm_enum.py
@@ -326,9 +326,9 @@ def main(args):
     kernel = {"nuts": NUTS, "hmc": HMC}[args.kernel](model)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
-        args.num_chains,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
+        num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
     mcmc.run(rng_key, sequences, lengths, args=args)
diff --git a/examples/neutra.py b/examples/neutra.py
@@ -64,8 +64,8 @@ def main(args):
     nuts_kernel = NUTS(dual_moon_model)
     mcmc = MCMC(
         nuts_kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
@@ -91,8 +91,8 @@ def main(args):
     nuts_kernel = NUTS(neutra_model)
     mcmc = MCMC(
         nuts_kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/ode.py b/examples/ode.py
@@ -90,8 +90,8 @@ def main(args):
     # use dense_mass for better mixing rate
     mcmc = MCMC(
         NUTS(model, dense_mass=True),
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/proportion_test.py b/examples/proportion_test.py
@@ -131,9 +131,9 @@ def run_inference(
     kernel = NUTS(model)
     mcmc = MCMC(
         kernel,
-        num_warmup,
-        num_samples,
-        num_chains,
+        num_warmup=num_warmup,
+        num_samples=num_samples,
+        num_chains=num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
     mcmc.run(rng_key, design_matrix, outcome)
diff --git a/examples/sparse_regression.py b/examples/sparse_regression.py
@@ -230,8 +230,8 @@ def run_inference(model, args, rng_key, X, Y, hypers):
     kernel = NUTS(model)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
         num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
diff --git a/examples/ucbadmit.py b/examples/ucbadmit.py
@@ -92,9 +92,9 @@ def run_inference(dept, male, applications, admit, rng_key, args):
     kernel = NUTS(glmm)
     mcmc = MCMC(
         kernel,
-        args.num_warmup,
-        args.num_samples,
-        args.num_chains,
+        num_warmup=args.num_warmup,
+        num_samples=args.num_samples,
+        num_chains=args.num_chains,
         progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
     )
     mcmc.run(rng_key, dept, male, applications, admit)
diff --git a/notebooks/source/bayesian_hierarchical_linear_regression.ipynb b/notebooks/source/bayesian_hierarchical_linear_regression.ipynb
@@ -678,7 +678,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/bayesian_imputation.ipynb b/notebooks/source/bayesian_imputation.ipynb
@@ -780,7 +780,7 @@
     }
    ],
    "source": [
-    "mcmc = MCMC(NUTS(model), 1000, 1000)\n",
+    "mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=1000)\n",
     "mcmc.run(random.PRNGKey(0), **data, survived=survived)\n",
     "mcmc.print_summary()"
    ]
@@ -1103,7 +1103,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/bayesian_regression.ipynb b/notebooks/source/bayesian_regression.ipynb
@@ -1269,11 +1269,9 @@
     "rng_key = random.PRNGKey(0)\n",
     "rng_key, rng_key_ = random.split(rng_key)\n",
     "\n",
-    "num_warmup, num_samples = 1000, 2000\n",
-    "\n",
     "# Run NUTS.\n",
     "kernel = NUTS(model)\n",
-    "mcmc = MCMC(kernel, num_warmup, num_samples)\n",
+    "mcmc = MCMC(kernel, num_warmup=1000, num_samples=2000)\n",
     "mcmc.run(rng_key_, marriage=dset.MarriageScaled.values, divorce=dset.DivorceScaled.values)\n",
     "mcmc.print_summary()\n",
     "samples_1 = mcmc.get_samples()"
@@ -2205,7 +2203,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/discrete_imputation.ipynb b/notebooks/source/discrete_imputation.ipynb
@@ -802,7 +802,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/logistic_regression.ipynb b/notebooks/source/logistic_regression.ipynb
@@ -373,7 +373,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/model_rendering.ipynb b/notebooks/source/model_rendering.ipynb
@@ -563,7 +563,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/ordinal_regression.ipynb b/notebooks/source/ordinal_regression.ipynb
@@ -281,7 +281,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/source/time_series_forecasting.ipynb b/notebooks/source/time_series_forecasting.ipynb
@@ -512,7 +512,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.8.8"
   }
  },
  "nbformat": 4,
diff --git a/numpyro/compat/infer.py b/numpyro/compat/infer.py
@@ -82,7 +82,7 @@ def __init__(
         self,
         kernel,
         num_samples,
-        warmup_steps=None,
+        num_warmup=None,
         initial_params=None,
         num_chains=1,
         hook_fn=None,
@@ -91,12 +91,12 @@ def __init__(
         disable_validation=True,
         transforms=None,
     ):
-        if warmup_steps is None:
-            warmup_steps = num_samples
+        if num_warmup is None:
+            num_warmup = num_samples
         self._initial_params = initial_params
         self._mcmc = mcmc.MCMC(
             kernel,
-            warmup_steps,
+            num_warmup,
             num_samples,
             num_chains=num_chains,
             progress_bar=(not disable_progbar),
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
@@ -47,7 +47,7 @@
    >>> labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
 
    >>> num_warmup, num_samples = 1000, 1000
-   >>> mcmc = MCMC(NUTS(model=logistic_regression), num_warmup, num_samples)
+   >>> mcmc = MCMC(NUTS(model=logistic_regression), num_warmup=num_warmup, num_samples=num_samples)
    >>> mcmc.run(random.PRNGKey(2), data, labels)  # doctest: +SKIP
    sample: 100%|██████████| 1000/1000 [00:00<00:00, 1252.39it/s, 1 steps of size 5.83e-01. acc. prob=0.85]
    >>> mcmc.print_summary()  # doctest: +SKIP
diff --git a/numpyro/infer/hmc_gibbs.py b/numpyro/infer/hmc_gibbs.py
@@ -86,7 +86,7 @@ class HMCGibbs(MCMCKernel):
         ...
         >>> hmc_kernel = NUTS(model)
         >>> kernel = HMCGibbs(hmc_kernel, gibbs_fn=gibbs_fn, gibbs_sites=['x'])
-        >>> mcmc = MCMC(kernel, 100, 100, progress_bar=False)
+        >>> mcmc = MCMC(kernel, num_warmup=100, num_samples=100, progress_bar=False)
         >>> mcmc.run(random.PRNGKey(0))
         >>> mcmc.print_summary()  # doctest: +SKIP
 
@@ -386,7 +386,7 @@ class DiscreteHMCGibbs(HMCGibbs):
         >>> probs = jnp.array([0.15, 0.3, 0.3, 0.25])
         >>> locs = jnp.array([-2, 0, 2, 4])
         >>> kernel = DiscreteHMCGibbs(NUTS(model), modified=True)
-        >>> mcmc = MCMC(kernel, 1000, 100000, progress_bar=False)
+        >>> mcmc = MCMC(kernel, num_warmup=1000, num_samples=100000, progress_bar=False)
         >>> mcmc.run(random.PRNGKey(0), probs, locs)
         >>> mcmc.print_summary()  # doctest: +SKIP
         >>> samples = mcmc.get_samples()["x"]
@@ -587,7 +587,7 @@ class HMCECS(HMCGibbs):
         ...
         >>> data = random.normal(random.PRNGKey(0), (10000,)) + 1
         >>> kernel = HMCECS(NUTS(model), num_blocks=10)
-        >>> mcmc = MCMC(kernel, 1000, 1000)
+        >>> mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
         >>> mcmc.run(random.PRNGKey(0), data)
         >>> samples = mcmc.get_samples()["x"]
         >>> assert abs(jnp.mean(samples) - 1.) < 0.1
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
@@ -234,6 +234,7 @@ class MCMC(object):
     def __init__(
         self,
         sampler,
+        *,
         num_warmup,
         num_samples,
         num_chains=1,
@@ -439,7 +440,7 @@ def post_warmup_state(self):
 
             .. code-block:: python
 
-                mcmc = MCMC(NUTS(model), 100, 100)
+                mcmc = MCMC(NUTS(model), num_warmup=100, num_samples=100)
                 mcmc.run(random.PRNGKey(0))
                 first_100_samples = mcmc.get_samples()
                 mcmc.post_warmup_state = mcmc.last_state
@@ -466,7 +467,7 @@ def warmup(
         extra_fields=(),
         collect_warmup=False,
         init_params=None,
-        **kwargs
+        **kwargs,
     ):
         """
         Run the MCMC warmup adaptation phase. After this call, `self.warmup_state` will be set
diff --git a/numpyro/infer/mixed_hmc.py b/numpyro/infer/mixed_hmc.py
@@ -59,7 +59,7 @@ class MixedHMC(DiscreteHMCGibbs):
         >>> probs = jnp.array([0.15, 0.3, 0.3, 0.25])
         >>> locs = jnp.array([-2, 0, 2, 4])
         >>> kernel = MixedHMC(HMC(model, trajectory_length=1.2), num_discrete_updates=20)
-        >>> mcmc = MCMC(kernel, 1000, 100000, progress_bar=False)
+        >>> mcmc = MCMC(kernel, num_warmup=1000, num_samples=100000, progress_bar=False)
         >>> mcmc.run(random.PRNGKey(0), probs, locs)
         >>> mcmc.print_summary()  # doctest: +SKIP
         >>> samples = mcmc.get_samples()
diff --git a/test/contrib/test_control_flow.py b/test/contrib/test_control_flow.py
@@ -36,7 +36,7 @@ def transition(state, i):
     T = 10
     num_samples = 100
     kernel = NUTS(model)
-    mcmc = MCMC(kernel, 100, num_samples)
+    mcmc = MCMC(kernel, num_warmup=100, num_samples=num_samples)
     mcmc.run(random.PRNGKey(0), T=T)
     assert set(mcmc.get_samples()) == {"x", "y", "y2", "x_0", "y_0"}
     mcmc.print_summary()
diff --git a/test/contrib/test_infer_discrete.py b/test/contrib/test_infer_discrete.py
@@ -360,7 +360,7 @@ def model2():
 @pytest.mark.parametrize("model", [model_zzxx, model2])
 @pytest.mark.parametrize("temperature", [0, 1])
 def test_mcmc_model_side_enumeration(model, temperature):
-    mcmc = infer.MCMC(infer.NUTS(model), 0, 1)
+    mcmc = infer.MCMC(infer.NUTS(model), num_warmup=0, num_samples=1)
     mcmc.run(random.PRNGKey(0))
     mcmc_data = {
         k: v[0] for k, v in mcmc.get_samples().items() if k in ["loc", "scale"]
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
@@ -169,7 +169,7 @@ def test_random_module__mcmc(backend, init):
         kwargs_name = "x"
 
     N, dim = 3000, 3
-    warmup_steps, num_samples = (1000, 1000)
+    num_warmup, num_samples = (1000, 1000)
     data = random.normal(random.PRNGKey(0), (N, dim))
     true_coefs = np.arange(1.0, dim + 1.0)
     logits = np.sum(true_coefs * data, axis=-1)
@@ -191,7 +191,9 @@ def model(data, labels):
         numpyro.sample("y", dist.Bernoulli(logits=logits), obs=labels)
 
     kernel = NUTS(model=model)
-    mcmc = MCMC(kernel, warmup_steps, num_samples, progress_bar=False)
+    mcmc = MCMC(
+        kernel, num_warmup=num_warmup, num_samples=num_samples, progress_bar=False
+    )
     mcmc.run(random.PRNGKey(2), data, labels)
     mcmc.print_summary()
     samples = mcmc.get_samples()
diff --git a/test/contrib/test_tfp.py b/test/contrib/test_tfp.py
diff --git a/test/infer/test_hmc_gibbs.py b/test/infer/test_hmc_gibbs.py
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
diff --git a/test/test_compile.py b/test/test_compile.py
diff --git a/test/test_pickle.py b/test/test_pickle.py

Original file line number	Diff line number	Diff line change
`@@ -192,7 +192,7 @@ def benchmark_hmc(args, features, labels):`
`192`	`192`	`)`
`193`	`193`	`else:`
`194`	`194`	`raise ValueError("Invalid algorithm, either 'HMC', 'NUTS', or 'HMCECS'.")`
`195`		`- mcmc = MCMC(kernel, args.num_warmup, args.num_samples)`
	`195`	`+ mcmc = MCMC(kernel, num_warmup=args.num_warmup, num_samples=args.num_samples)`
`196`	`196`	`mcmc.run(rng_key, features, labels, subsample_size, extra_fields=("accept_prob",))`
`197`	`197`	`print("Mean accept prob:", jnp.mean(mcmc.get_extra_fields()["accept_prob"]))`
`198`	`198`	`mcmc.print_summary(exclude_deterministic=False)`