Fix observed variable detection

scvi/model/base/_pyromixin.py

@@ -160,6 +160,7 @@ def _get_one_posterior_sample(
         kwargs,
         return_sites: Optional[list] = None,
         return_observed: bool = False,
+        observed_not_deterministic: Optional[list] = None,
     ):
         """
         Get one sample from posterior distribution.
@@ -174,6 +175,8 @@ def _get_one_posterior_sample(
             List of variables for which to generate posterior samples, defaults to all variables.
         return_observed
             Record samples of observed variables.
+        observed_not_deterministic
+            List of observed but non-deterministic variables.
 
         Returns
         -------
@@ -195,18 +198,17 @@ def _get_one_posterior_sample(
                     and (
                         (return_sites is None) or (name in return_sites)
                     )  # selected in return_sites list
-                    and (
-                        (
-                            (not site.get("is_observed", True)) or return_observed
-                        )  # don't save observed unless requested
-                        or (site.get("infer", False).get("_deterministic", False))
-                    )  # unless it is deterministic
                     and not isinstance(
                         site.get("fn", None), poutine.subsample_messenger._Subsample
                     )  # don't save plates
                 )
             }
 
+        if not return_observed:
+            sample = {
+                k: v for k, v in sample.items() if k not in observed_not_deterministic
+            }
+
         sample = {name: site.cpu().numpy() for name, site in sample.items()}
 
         return sample
@@ -241,8 +243,13 @@ def _get_posterior_samples(
         Dictionary with array of samples for each variable
         dictionary {variable_name: [array with samples in 0 dimension]}
         """
+        observed_not_deterministic = self._get_observed_sites(args, kwargs)
         samples = self._get_one_posterior_sample(
-            args, kwargs, return_sites=return_sites, return_observed=return_observed
+            args,
+            kwargs,
+            return_sites=return_sites,
+            return_observed=return_observed,
+            observed_not_deterministic=observed_not_deterministic,
         )
         samples = {k: [v] for k, v in samples.items()}
 
@@ -255,7 +262,11 @@ def _get_posterior_samples(
 
             # generate new sample
             samples_ = self._get_one_posterior_sample(
-                args, kwargs, return_sites=return_sites, return_observed=return_observed
+                args,
+                kwargs,
+                return_sites=return_sites,
+                return_observed=return_observed,
+                observed_not_deterministic=observed_not_deterministic,
             )
 
             # add new sample
@@ -309,21 +320,58 @@ def _get_obs_plate_sites(
             for name, site in trace.nodes.items()
             if (
                 (site["type"] == "sample")  # sample statement
-                and (
-                    (
-                        (not site.get("is_observed", True)) or return_observed
-                    )  # don't save observed unless requested
-                    or (site.get("infer", False).get("_deterministic", False))
-                )  # unless it is deterministic
                 and not isinstance(
                     site.get("fn", None), poutine.subsample_messenger._Subsample
                 )  # don't save plates
             )
             if any(f.name == plate_name for f in site["cond_indep_stack"])
         }
+        if not return_observed:
+            observed_not_deterministic = self._get_observed_sites(args, kwargs)
+            obs_plate = {
+                k: v
+                for k, v in obs_plate.items()
+                if k not in observed_not_deterministic
+            }
 
         return obs_plate
 
+    def _get_observed_sites(
+        self,
+        args: list,
+        kwargs: dict,
+    ):
+        """
+        Automatically guess which model sites correspond to observed variables
+
+        This excludes pyro.deterministic variables.
+
+        Parameters
+        ----------
+        args
+            Arguments to the model.
+        kwargs
+            Keyword arguments to the model.
+
+        Returns
+        -------
+        List with site names.
+        """
+        trace = poutine.trace(self.module.model).get_trace(*args, **kwargs)
+        observed_sites = [
+            name
+            for name, site in trace.nodes.items()
+            if (
+                (site["type"] == "sample")  # sample statement
+                and (
+                    site.get("is_observed", True)
+                    and site.get("infer", False).get("_deterministic", False)
+                )  # exclude deterministic sites
+            )
+        ]
+
+        return observed_sites
+
     def _posterior_samples_minibatch(
         self, use_gpu: bool = None, batch_size: Optional[int] = None, **sample_kwargs
     ):

tests/models/test_pyro.py

@@ -100,6 +100,7 @@ def forward(self, x, y, ind_x):
                     "per_cell_weights", dist.Normal(self.zero, self.one)
                 )
                 mean = mean + per_cell_weights.squeeze(-1)
+                pyro.deterministic("weighted_mean", mean)
 
         with obs_plate:
             pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
@@ -396,6 +397,10 @@ def test_pyro_bayesian_train_sample_mixin_with_local():
         adata.n_obs,
         1,
     )
+    # test that observed variables are excluded
+    assert "obs" not in samples["posterior_samples"].keys()
+    # test that deterministic variables are included
+    assert "weighted_mean" in samples["posterior_samples"].keys()
 
 
 def test_pyro_bayesian_train_sample_mixin_with_local_full_data():
@@ -614,6 +619,14 @@ def test_lda_model():
     mod.get_elbo(adata2)
     mod.get_perplexity(adata2)
 
+    # test posterior sampling
+    samples = mod.sample_posterior(
+        num_samples=10, use_gpu=use_gpu, batch_size=adata.n_obs, return_samples=True
+    )
+    assert samples["posterior_samples"]["latent"].shape == (10, adata.n_obs, n_topics)
+    # test that observed variables are excluded
+    assert "obs" not in samples["posterior_samples"].keys()
+
 
 def test_lda_model_save_load(save_path):
     use_gpu = torch.cuda.is_available()