fix returned object when no vars to sample and extend=True

pymc-devs · Sep 18, 2024 · b9fbfed · b9fbfed
1 parent 4300be1
commit b9fbfed
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/pymc/sampling/forward.py b/pymc/sampling/forward.py
@@ -825,7 +825,7 @@ def sample_posterior_predictive(
         if return_inferencedata and not extend_inferencedata:
             return InferenceData()
         elif return_inferencedata and extend_inferencedata:
-            return trace
+            return trace if idata is None else idata
         return {}
 
     vars_in_trace = get_vars_in_point_list(_trace, model)

diff --git a/tests/sampling/test_forward.py b/tests/sampling/test_forward.py
@@ -492,6 +492,11 @@ def test_normal_scalar(self):
             ppc = pm.sample_posterior_predictive(trace, var_names=[], return_inferencedata=False)
             assert len(ppc) == 0
 
+            # test empty ppc with extend_inferencedata
+            assert isinstance(trace, InferenceData)
+            ppc = pm.sample_posterior_predictive(trace, var_names=[], extend_inferencedata=True)
+            assert ppc is trace
+
             # test keep_size parameter
             ppc = pm.sample_posterior_predictive(trace, return_inferencedata=False)
             assert ppc["a"].shape == (nchains, ndraws)