Check for observed variables in the trace as well as the model

Author: zaxtax

pymc/sampling/forward.py

@@ -767,13 +767,15 @@ def sample_posterior_predictive(
     if "coords" not in idata_kwargs:
         idata_kwargs["coords"] = {}
     idata: InferenceData | None = None
+    observed_data = None
     stacked_dims = None
     if isinstance(trace, InferenceData):
         _constant_data = getattr(trace, "constant_data", None)
         if _constant_data is not None:
             trace_coords.update({str(k): v.data for k, v in _constant_data.coords.items()})
             constant_data.update({str(k): v.data for k, v in _constant_data.items()})
         idata = trace
+        observed_data = trace["observed_data"]
         trace = trace["posterior"]
     if isinstance(trace, xarray.Dataset):
         trace_coords.update({str(k): v.data for k, v in trace.coords.items()})
@@ -817,6 +819,8 @@ def sample_posterior_predictive(
         vars_ = [model[x] for x in var_names]
     else:
         vars_ = model.observed_RVs + observed_dependent_deterministics(model)
+        if observed_data is not None:
+            vars_ += [model[x] for x in observed_data if x in model]
 
     vars_to_sample = list(get_default_varnames(vars_, include_transformed=False))
 

tests/sampling/test_forward.py

@@ -481,6 +481,16 @@ def test_normal_scalar(self):
                 chains=nchains,
             )
 
+        # test that trace is used in ppc
+        with pm.Model() as model_ppc:
+            mu = pm.Normal("mu", 0.0, 1.0)
+            a = pm.Normal("a", mu=mu, sigma=1)
+
+        ppc = pm.sample_posterior_predictive(
+            trace=trace, model=model_ppc, return_inferencedata=False
+        )
+        assert "a" in ppc
+
         with model:
             # test list input
             ppc0 = pm.sample_posterior_predictive(