DOC Add Predictive examples

numpyro/infer/mcmc.py

@@ -64,7 +64,9 @@ class MCMCKernel(ABC):
         >>> kernel = MetropolisHastings(f)
         >>> mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
         >>> mcmc.run(random.PRNGKey(0), init_params=jnp.array([1., 2.]))
-        >>> samples = mcmc.get_samples()
+        >>> posterior_samples = mcmc.get_samples()
+        >>> predictive = Predictive(model, posterior_samples=posterior_samples)
+        >>> samples = predictive(rng_key1, *model_args, **model_kwargs)
         >>> mcmc.print_summary()  # doctest: +SKIP
     """
 
@@ -509,6 +511,8 @@ def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
         """
         Run the MCMC samplers and collect samples.
 
+        See :class:`~numpyro.infer.util.Predictive` for how to use them to collect posterior predictive samples.
+
         :param random.PRNGKey rng_key: Random number generator key to be used for the sampling.
             For multi-chains, a batch of `num_chains` keys can be supplied. If `rng_key`
             does not have batch_size, it will be split in to a batch of `num_chains` keys.

numpyro/infer/svi.py

@@ -83,7 +83,7 @@ class SVI(object):
         >>> import numpyro
         >>> import numpyro.distributions as dist
         >>> from numpyro.distributions import constraints
-        >>> from numpyro.infer import SVI, Trace_ELBO
+        >>> from numpyro.infer import Predictive, SVI, Trace_ELBO
 
         >>> def model(data):
         ...     f = numpyro.sample("latent_fairness", dist.Beta(10, 10))
@@ -101,6 +101,8 @@ class SVI(object):
         >>> svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
         >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)
         >>> params = svi_result.params
+        >>> predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
+        >>> samples = predictive(random.PRNGKey(0), data)
         >>> inferred_mean = params["alpha_q"] / (params["alpha_q"] + params["beta_q"])
 
     :param model: Python callable with Pyro primitives for the model.

numpyro/infer/util.py

@@ -743,6 +743,21 @@ class Predictive(object):
           argument is not None, its value should be equal to `num_chains x N`.
 
     :return: dict of samples from the predictive distribution.
+
+    **Example:**
+
+    Given a model, you can sample from the prior predictive:
+
+        >>> predictive = Predictive(model, num_samples=num_samples)
+        >>> samples = predictive(rng_key1, *model_args, **model_kwargs)
+
+    If you also have posterior samples, you can sample from the posterior predictive:
+
+        >>> predictive = Predictive(model, posterior_samples=posterior_samples)
+        >>> samples = predictive(rng_key1, *model_args, **model_kwargs)
+
+    See docstrings for :class:`~numpyro.infer.svi.SVI` and :class:`~numpyro.infer.mcmc.MCMCKernel`
+    to see example code of this in context.
     """
 
     def __init__(