DOC Add Predictive examples

numpyro/infer/mcmc.py

@@ -64,7 +64,7 @@ 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()
         >>> mcmc.print_summary()  # doctest: +SKIP
     """
 
@@ -593,6 +593,15 @@ def get_samples(self, group_by_chain=False):
             `dict` keyed on site names if a model containing Pyro primitives is used,
             but can be any :func:`jaxlib.pytree`, more generally (e.g. when defining a
             `potential_fn` for HMC that takes `list` args).
+
+        **Example:**
+
+        You can then pass those samples to :class:`~numpyro.infer.util.Predictive`::
+
+            posterior_samples = mcmc.get_samples()
+            predictive = Predictive(model, posterior_samples=posterior_samples)
+            samples = predictive(rng_key1, *model_args, **model_kwargs)
+
         """
         return (
             self._states[self._sample_field]

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))
@@ -102,6 +102,9 @@ class SVI(object):
         >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)
         >>> params = svi_result.params
         >>> inferred_mean = params["alpha_q"] / (params["alpha_q"] + params["beta_q"])
+        >>> # get posterior samples
+        >>> predictive = Predictive(guide, params=params, num_samples=1000)
+        >>> samples = predictive(random.PRNGKey(1), data)
 
     :param model: Python callable with Pyro primitives for the model.
     :param guide: Python callable with Pyro primitives for the guide

numpyro/infer/util.py

@@ -778,6 +778,27 @@ 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
+
+        def model(X, y=None):
+            ...
+            return numpyro.sample("obs", likelihood, obs=y)
+
+    you can sample from the prior predictive:
+
+        predictive = Predictive(model, num_samples=1000)
+        y_pred = predictive(rng_key, X)["obs"]
+
+    If you also have posterior samples, you can sample from the posterior predictive:
+
+        predictive = Predictive(model, posterior_samples=posterior_samples)
+        y_pred = predictive(rng_key, X)["obs"]
+
+    See docstrings for :class:`~numpyro.infer.svi.SVI` and :class:`~numpyro.infer.mcmc.MCMCKernel`
+    to see example code of this in context.
     """
 
     def __init__(