Predictive: better docs for setting obs to None

[tomwallis]

Hello,

It wasn't clear to me until I dug into some strange behaviour that the Predictive method requires setting obs to None. I see now that the examples do it, but it's not made explicit that this is required for the class to work as intended. Could this be made more explicit in the docs? Thanks in advance!

[fehiepsi]

You are right. It is not trivial to realize that api. Do you want to make a PR to enhance the docs?

[tomwallis]

I can try. To do this, should I build the docs locally, or are the changes to the source files sufficient?

[tillahoffmann]

Expanding the documentation in the source file linked below should do the trick. You don't need to build the documentation locally, but it can be helpful for checking that the output looks like what you intended.

numpyro/numpyro/infer/util.py

Lines 851 to 918 in d61f15c

	class Predictive(object):
	"""
	This class is used to construct predictive distribution. The predictive distribution is obtained
	by running model conditioned on latent samples from `posterior_samples`.
	.. warning::
	The interface for the `Predictive` class is experimental, and
	might change in the future.
	:param model: Python callable containing Pyro primitives.
	:param dict posterior_samples: dictionary of samples from the posterior.
	:param callable guide: optional guide to get posterior samples of sites not present
	in `posterior_samples`.
	:param dict params: dictionary of values for param sites of model/guide.
	:param int num_samples: number of samples
	:param list return_sites: sites to return; by default only sample sites not present
	in `posterior_samples` are returned.
	:param bool infer_discrete: whether or not to sample discrete sites from the
	posterior, conditioned on observations and other latent values in
	``posterior_samples``. Under the hood, those sites will be marked with
	``site["infer"]["enumerate"] = "parallel"``. See how `infer_discrete` works at
	the `Pyro enumeration tutorial <https://pyro.ai/examples/enumeration.html>`_.
	Note that this requires ``funsor`` installation.
	:param bool parallel: whether to predict in parallel using JAX vectorized map :func:`jax.vmap`.
	Defaults to False.
	:param batch_ndims: the number of batch dimensions in posterior samples or parameters. If `None` defaults
	to 0 if guide is set (i.e. not `None`) and 1 otherwise. Usages for batched posterior samples:
	+ set `batch_ndims=0` to get prediction for 1 single sample
	+ set `batch_ndims=1` to get prediction for `posterior_samples`
	with shapes `(num_samples x ...)` (same as`batch_ndims=None` with `guide=None`)
	+ set `batch_ndims=2` to get prediction for `posterior_samples`
	with shapes `(num_chains x N x ...)`. Note that if `num_samples`
	argument is not None, its value should be equal to `num_chains x N`.
	Usages for batched parameters:
	+ set `batch_ndims=0` to get 1 sample from the guide and parameters (same as `batch_ndims=None` with guide)
	+ set `batch_ndims=1` to get predictions from a one dimensional batch of the guide and parameters
	with shapes `(num_samples x batch_size x ...)`
	:param exclude_deterministic: indicates whether to ignore deterministic sites from the posterior samples.
	: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.
	"""

[tomwallis]

Thanks @tillahoffmann. I have created a PR that hopefully does the trick. I was not able to build the docs locally without more time investment, so I have not checked the intended output.