Repository: pyro-ppl/pyro

Issue:
[discussion] Keeping subsample correlations in TracePredictive
While developing #1854, I noticed that we do new random subsampling for each stochastic site in `TracePredictive` distribution. After reflecting about this, I have figured out that this might drop correlations between the sampled parameters under a subsampling statement.

I think what one needs to do is to fix a subset of random indices for each subsample site for all the sites. Any ideas on how one could this in a nice fashion?

My current idea is to create a dictionary when fixing the subsample sites, that map each subsample site to the set of random indices based on the provided subsample site size. Would this work?

Thank you in advance! 😄 


Discussion:

Resolving patch:
diff --git a/pyro/infer/abstract_infer.py b/pyro/infer/abstract_infer.py
--- a/pyro/infer/abstract_infer.py
+++ b/pyro/infer/abstract_infer.py
@@ -313,25 +313,28 @@ def _remove_dropped_nodes(self, trace):
                 continue
 
     def _adjust_to_data(self, trace, data_trace):
-        for name, site in list(trace.nodes.items()):
+        subsampled_idxs = dict()
+        for name, site in trace.iter_stochastic_nodes():
             # Adjust subsample sites
             if site_is_subsample(site):
                 site["fn"] = data_trace.nodes[name]["fn"]
                 site["value"] = data_trace.nodes[name]["value"]
             # Adjust sites under conditionally independent stacks
-            try:
-                site["cond_indep_stack"] = data_trace.nodes[name]["cond_indep_stack"]
-                site["fn"] = data_trace.nodes[name]["fn"]
-                for cis in site["cond_indep_stack"]:
-                    # Select random sub-indices to replay values under conditionally independent stacks.
-                    # Otherwise, we assume there is an dependence of indexes between training data
-                    # and prediction data.
-                    batch_dim = cis.dim - site["fn"].event_dim
-                    subidxs = torch.randint(0, site['value'].size(batch_dim), (cis.size,),
-                                            device=site["value"].device)
-                    site["value"] = site["value"].index_select(batch_dim, subidxs)
-            except KeyError:
-                pass
+            orig_cis_stack = site["cond_indep_stack"]
+            site["cond_indep_stack"] = data_trace.nodes[name]["cond_indep_stack"]
+            assert len(orig_cis_stack) == len(site["cond_indep_stack"])
+            site["fn"] = data_trace.nodes[name]["fn"]
+            for ocis, cis in zip(orig_cis_stack, site["cond_indep_stack"]):
+                # Select random sub-indices to replay values under conditionally independent stacks.
+                # Otherwise, we assume there is an dependence of indexes between training data
+                # and prediction data.
+                assert ocis.name == cis.name
+                assert not site_is_subsample(site)
+                batch_dim = cis.dim - site["fn"].event_dim
+                subsampled_idxs[cis.name] = subsampled_idxs.get(cis.name,
+                                                                torch.randint(0, ocis.size, (cis.size,),
+                                                                              device=site["value"].device))
+                site["value"] = site["value"].index_select(batch_dim, subsampled_idxs[cis.name])
 
     def marginal(self, sites=None):
         """


Test patch:
diff --git a/tests/infer/test_abstract_infer.py b/tests/infer/test_abstract_infer.py
--- a/tests/infer/test_abstract_infer.py
+++ b/tests/infer/test_abstract_infer.py
@@ -19,6 +19,13 @@ def model(num_trials):
         return pyro.sample("obs", dist.Binomial(num_trials, success_prob))
 
 
+def one_hot_model(pseudocounts, classes=None):
+    probs_prior = dist.Dirichlet(pseudocounts)
+    probs = pyro.sample("probs", probs_prior)
+    with pyro.plate("classes", classes.size(0) if classes is not None else 1, dim=-1):
+        return pyro.sample("obs", dist.OneHotCategorical(probs), obs=classes)
+
+
 def beta_guide(num_trials):
     phi_c0 = pyro.param("phi_c0", num_trials.new_tensor(5.0).expand([num_trials.size(0)]))
     phi_c1 = pyro.param("phi_c1", num_trials.new_tensor(5.0).expand([num_trials.size(0)]))
@@ -81,6 +88,19 @@ def test_posterior_predictive_svi_auto_diag_normal_guide():
     assert_close(marginal_return_vals.mean, torch.ones(5) * 700, rtol=0.05)
 
 
+def test_posterior_predictive_svi_one_hot():
+    pseudocounts = torch.ones(3) * 0.1
+    true_probs = torch.tensor([0.15, 0.6, 0.25])
+    classes = dist.OneHotCategorical(true_probs).sample((10000,))
+    opt = optim.Adam(dict(lr=0.1))
+    loss = Trace_ELBO()
+    guide = AutoDelta(one_hot_model)
+    svi_run = SVI(one_hot_model, guide, opt, loss, num_steps=1000, num_samples=1000).run(pseudocounts, classes=classes)
+    posterior_predictive = TracePredictive(one_hot_model, svi_run, num_samples=10000).run(pseudocounts)
+    marginal_return_vals = posterior_predictive.marginal().empirical["_RETURN"]
+    assert_close(marginal_return_vals.mean, true_probs.unsqueeze(0), rtol=0.1)
+
+
 def test_posterior_predictive_svi_auto_delta_guide_large_eval():
     true_probs = torch.ones(5) * 0.7
     num_trials = torch.ones(5) * 1000