Change include_samples to return_samples for more consistent API

Author: zaxtax

pymc_experimental/model/marginal_model.py

@@ -279,7 +279,7 @@ def unmarginalize(self, rvs_to_unmarginalize):
             self.register_rv(rv, name=rv.name)
 
     def recover_marginals(
-        self, idata, var_names=None, include_samples=False, extend_inferencedata=True
+        self, idata, var_names=None, return_samples=False, extend_inferencedata=True
     ):
         """Computes unnormalized posterior probabilities of marginalized variables
         conditioned on parameters of the model given InferenceData with posterior group
@@ -295,8 +295,8 @@ def recover_marginals(
             InferenceData with posterior group
         var_names : sequence of str, optional
             List of Observed variable names for which to compute log_likelihood. Defaults to all observed variables
-        include_samples : bool, default False
-            Include samples of the marginalized variables
+        return_samples : bool, default False
+            If True, also return samples of the marginalized variables
         extend_inferencedata : bool, default True
             Whether to extend the original InferenceData or return a new one
 
@@ -360,7 +360,7 @@ def transform_input(inputs):
             )
 
             rv_loglike_fn = None
-            if include_samples:
+            if return_samples:
                 sample_rv_outs = pymc.Categorical.dist(logit_p=joint_logps)
                 rv_loglike_fn = compile_pymc(
                     inputs=other_values,
@@ -376,7 +376,7 @@ def transform_input(inputs):
 
             logvs = [rv_loglike_fn(**vs) for vs in posterior_pts]
 
-            if include_samples:
+            if return_samples:
                 logps, samples = zip(*logvs)
                 logps = np.array(logps)
                 rv_dict[rv.name] = np.reshape(

pymc_experimental/tests/model/test_marginal_model.py

@@ -274,7 +274,7 @@ def test_recover_marginals_basic():
         )
         idata = InferenceData(posterior=dict_to_dataset(prior))
 
-    idata = m.recover_marginals(idata, include_samples=True)
+    idata = m.recover_marginals(idata, return_samples=True)
     assert "k" in idata
     assert "lp_k" in idata
     assert idata.k.shape == idata.y.shape
@@ -313,7 +313,7 @@ def test_nested_recover_marginals():
         )
         idata = InferenceData(posterior=dict_to_dataset(prior))
 
-    idata = m.recover_marginals(idata, include_samples=True)
+    idata = m.recover_marginals(idata, return_samples=True)
     assert "idx" in idata
     assert "lp_idx" in idata
     assert idata.idx.shape == idata.y.shape