Update docs to make Predictive behavior more clear

notebooks/source/bayesian_hierarchical_linear_regression.ipynb

@@ -427,6 +427,13 @@
     "samples_predictive = predictive(random.PRNGKey(0), patient_code, Weeks, None)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that for [`Predictive`](http://num.pyro.ai/en/latest/utilities.html#numpyro.infer.util.Predictive) to work as expected, the response variable of the model (in this case, `FVC_obs`) must be set to `None`."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

notebooks/source/bayesian_regression.ipynb

@@ -1176,7 +1176,7 @@
     "    \"South\",\n",
     "    \"Divorce\",\n",
     "]\n",
-    "sns.pairplot(dset, x_vars=vars, y_vars=vars, palette=\"husl\");"
+    "sns.pairplot(dset, x_vars=vars, y_vars=vars, palette=\"husl\")"
    ]
   },
   {
@@ -1218,7 +1218,7 @@
     }
    ],
    "source": [
-    "sns.regplot(x=\"WaffleHouses\", y=\"Divorce\", data=dset);"
+    "sns.regplot(x=\"WaffleHouses\", y=\"Divorce\", data=dset)"
    ]
   },
   {
@@ -1450,7 +1450,7 @@
     "ax = plot_regression(dset.MarriageScaled.values, mean_mu, hpdi_mu)\n",
     "ax.set(\n",
     "    xlabel=\"Marriage rate\", ylabel=\"Divorce rate\", title=\"Regression line with 90% CI\"\n",
-    ");"
+    ")"
    ]
   },
   {
@@ -1504,7 +1504,15 @@
     "hpdi_prior_pred = hpdi(prior_predictions, 0.9)\n",
     "\n",
     "ax = plot_regression(dset.MarriageScaled.values, mean_prior_pred, hpdi_prior_pred)\n",
-    "ax.set(xlabel=\"Marriage rate\", ylabel=\"Divorce rate\", title=\"Predictions with 90% CI\");"
+    "ax.set(xlabel=\"Marriage rate\", ylabel=\"Divorce rate\", title=\"Predictions with 90% CI\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that for `Predictive` to work as expected, the response variable of the model (in this case, `divorce`) must be set to `None`.\n",
+    "In the code above this is done implicitly by not passing a value for `divorce` to the model in the call to `prior_predictive`, which due to the model definition, sets `divorce=None`."
    ]
   },
   {
@@ -1782,7 +1790,7 @@
     "hpdi_pred = hpdi(predictions_1, 0.9)\n",
     "\n",
     "ax = plot_regression(dset.MarriageScaled.values, mean_pred, hpdi_pred)\n",
-    "ax.set(xlabel=\"Marriage rate\", ylabel=\"Divorce rate\", title=\"Predictions with 90% CI\");"
+    "ax.set(xlabel=\"Marriage rate\", ylabel=\"Divorce rate\", title=\"Predictions with 90% CI\")"
    ]
   },
   {
@@ -1963,7 +1971,7 @@
     "    xlabel=\"Median marriage age\",\n",
     "    ylabel=\"Divorce rate\",\n",
     "    title=\"Regression line with 90% CI\",\n",
-    ");"
+    ")"
    ]
   },
   {
@@ -2000,7 +2008,7 @@
     "hpdi_pred = hpdi(predictions_2, 0.9)\n",
     "\n",
     "ax = plot_regression(dset.AgeScaled.values, mean_pred, hpdi_pred)\n",
-    "ax.set(xlabel=\"Median Age\", ylabel=\"Divorce rate\", title=\"Predictions with 90% CI\");"
+    "ax.set(xlabel=\"Median Age\", ylabel=\"Divorce rate\", title=\"Predictions with 90% CI\")"
    ]
   },
   {
@@ -2216,7 +2224,7 @@
     ")\n",
     "ax[1].set(xlabel=\"Residuals\", ylabel=\"State\", title=\"Residuals with 90% CI\")\n",
     "ax[1].set_yticks(y)\n",
-    "ax[1].set_yticklabels(dset.Loc.values[idx], fontsize=10);"
+    "ax[1].set_yticklabels(dset.Loc.values[idx], fontsize=10)"
    ]
   },
   {
@@ -2468,7 +2476,7 @@
     "    \"Residuals (with error-bars) from current model (in red). \"\n",
     "    \"Black marker \\nshows residuals from the previous model (Model 3). \"\n",
     "    \"Measurement \\nerror is indicated by orange bar.\",\n",
-    ");"
+    ")"
    ]
   },
   {
@@ -2522,7 +2530,7 @@
     "    xlabel=\"Measurement Noise\",\n",
     "    ylabel=\"Residual\",\n",
     "    title=\"Mean residuals (Model 4: red, Model 3: blue)\",\n",
-    ");"
+    ")"
    ]
   },
   {

numpyro/infer/util.py

@@ -1,19 +1,18 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import warnings
 from collections import namedtuple
 from contextlib import contextmanager
 from functools import partial
 from typing import Callable, Optional
-import warnings
-
-import numpy as np
 
 import jax
+import jax.numpy as jnp
+import numpy as np
 from jax import device_get, jacfwd, lax, random, value_and_grad
 from jax.flatten_util import ravel_pytree
 from jax.lax import broadcast_shapes
-import jax.numpy as jnp
 
 import numpyro
 from numpyro.distributions import constraints
@@ -857,6 +856,9 @@ class Predictive(object):
         The interface for the `Predictive` class is experimental, and
         might change in the future.
 
+    Note that for the predictive distribution to be returned as intended, observed
+    variables in the model (constraining the likelihood term) must be set to `None` (see Example).
+
     :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
@@ -908,6 +910,10 @@ def model(X, y=None):
         predictive = Predictive(model, num_samples=1000)
         y_pred = predictive(rng_key, X)["obs"]
 
+    Note how above, no value for `y` is passed to `predictive`, resulting in `y`
+    being set to `None`. Setting the observed variable(s) to `None` when using
+    `Predictive` is required for the method to function as expected.
+
     If you also have posterior samples, you can sample from the posterior predictive::
 
         predictive = Predictive(model, posterior_samples=posterior_samples)