Repository: pyro-ppl/pyro

Issue:
map_data requires batch_size arg in both model and guide
To use data sub sampling (using `map_data`) we have to specify the `batch_size` arg in both model and guide, even though specifying this in the guide only is sufficient. (For SVI at least.)

(If no `batch_size` is given in the model, it appears that no sub-sampling happens (in the model). If a `batch_size` is given that doesn't match the `batch_size` in the guide then things appear to work, so I guess sub-sampling happens reusing indices from the guide, and `batch_size` is ignored.)

ETA: This also appears to be true of `iarange`.


Discussion:
What is the desired behavior? My guess is that [the short-circuiting logic](https://github.com/uber/pyro/blob/dev/pyro/__init__.py#L205) is preventing the model `map_data` from being replayed as you expect, but I'm unclear what the desired behavior is.
> What is the desired behavior?

As a user, when I want to enable sub sampling, it would be good to only have to specify the batch size in one place. (I think this makes sense because I never want batch sizes to differ across the model and guide, for now at least.) Personally, I suspect that batch size should eventually be given as a parameter of inference rather than via the model and/or guide (#236), but in the meantime, having to only specify it in the guide would be an improvement.

Right now we're in a slightly awkward situation where (for SVI at least) the batch size I give in the guide controls subsampling, but I have to pass *some* batch size to the model otherwise no sub sampling occurs, and inference dies because of the differing batch sizes. However, it does't matter *what* value I pass in the model, since the batch size from the guide controls the batch size. This is potentially confusing, because I might start tweaking the batch size in the model and see no effect.

I'm reopening this issue since this is still the case after #330 afaict.

(This isn't high priority for me.)

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -149,8 +149,13 @@ def sample(self):
         :returns: a random subsample of `range(size)`
         :rtype: torch.autograd.Variable of torch.LongTensor
         """
-        assert 0 <= self.subsample_size <= self.size
-        result = Variable(torch.randperm(self.size)[:self.subsample_size])
+        subsample_size = self.subsample_size
+        if subsample_size is None or subsample_size > self.size:
+            subsample_size = self.size
+        if subsample_size == self.size:
+            result = Variable(torch.LongTensor(list(range(self.size))))
+        else:
+            result = Variable(torch.randperm(self.size)[:self.subsample_size])
         return result.cuda() if self.use_cuda else result
 
     def batch_log_pdf(self, x):
@@ -202,25 +207,20 @@ def iarange(name, size=None, subsample_size=None, subsample=None, use_cuda=False
                 observe('obs', normal, data.index_select(0, ind), mu, sigma)
     """
     if size is None:
-        # Case: with iarange("name"): ...
         assert subsample_size is None
         assert subsample is None
         size = 1
         subsample_size = 1
-    elif subsample is not None:
-        # Case: with iarange("name", size, subsample=...): ...
-        subsample_size = len(subsample)
-        assert subsample_size <= size, 'subsample is larger than size'
-    elif subsample_size is None or subsample_size >= size:
-        # Case: with iarange("name", size) as ind: ...
-        subsample_size = size
-        subsample = Variable(torch.LongTensor(list(range(size))))
-        if use_cuda:
-            subsample = subsample.cuda()
-    else:
-        # Case: with iarange("name", size, subsample_size) as ind: ...
+    elif subsample is None:
         subsample = sample(name, _Subsample(size, subsample_size, use_cuda))
 
+    if subsample_size is None:
+        subsample_size = len(subsample)
+    elif subsample_size != len(subsample):
+        raise ValueError("subsample_size does not match len(subsample), {} vs {}.".format(
+            subsample_size, len(subsample)) +
+            " Did you accidentally use different subsample_size in the model and guide?")
+
     if len(_PYRO_STACK) == 0:
         yield subsample
     else:
@@ -249,12 +249,9 @@ def irange(name, size, subsample_size=None, subsample=None, use_cuda=False):
                 if z[i]:  # Prevents vectorization.
                     observe('obs_{}'.format(i), normal, data[i], mu, sigma)
     """
-    if subsample is not None:
-        subsample_size = len(subsample)
-
     with iarange(name, size, subsample_size, subsample, use_cuda) as ind:
         # Wrap computation in an independence context.
-        indep_context = LambdaPoutine(None, name, 1.0, 'list', 0, subsample_size)
+        indep_context = LambdaPoutine(None, name, 1.0, 'list', 0, len(ind))
         if isinstance(ind, Variable):
             ind = ind.data
         for i in ind:


Test patch:
diff --git a/tests/integration_tests/test_tracegraph_elbo.py b/tests/integration_tests/test_tracegraph_elbo.py
--- a/tests/integration_tests/test_tracegraph_elbo.py
+++ b/tests/integration_tests/test_tracegraph_elbo.py
@@ -591,10 +591,10 @@ def obs_outer(i, x):
         guide_trace = pyro.poutine.trace(guide, graph_type="dense").get_trace()
         model_trace = pyro.poutine.trace(pyro.poutine.replay(model, guide_trace),
                                          graph_type="dense").get_trace()
-        self.assertEqual(len(model_trace.edges()), 9)
-        self.assertEqual(len(model_trace.nodes()), 12)
-        self.assertEqual(len(guide_trace.edges()), 0)
-        self.assertEqual(len(guide_trace.nodes()), 5)
+        assert len(model_trace.edges()) == 34
+        assert len(model_trace.nodes()) == 16
+        assert len(guide_trace.edges()) == 7
+        assert len(guide_trace.nodes()) == 9
 
         adam = optim.Adam({"lr": 0.0008, "betas": (0.96, 0.999)})
         svi = SVI(model, guide, adam, loss="ELBO", trace_graph=True)
diff --git a/tests/poutine/test_mapdata.py b/tests/poutine/test_mapdata.py
--- a/tests/poutine/test_mapdata.py
+++ b/tests/poutine/test_mapdata.py
@@ -303,3 +303,31 @@ def test_cuda(model, subsample_size):
     tr = poutine.trace(model).get_trace(subsample_size)
     assert tr.log_pdf().is_cuda
     assert tr.batch_log_pdf().is_cuda
+
+
+@pytest.mark.parametrize('model', [
+    iarange_model,
+    irange_model,
+    map_data_vector_model,
+    map_data_iter_model,
+], ids=['iarange', 'irange', 'map_data_vector', 'map_data_iter'])
+@pytest.mark.parametrize("behavior,model_size,guide_size", [
+    ("error", 20, 5),
+    ("error", 5, 20),
+    ("error", 5, None),
+    ("ok", 20, 20),
+    ("ok", 20, None),
+    ("ok", 5, 5),
+    ("ok", None, 20),
+    ("ok", None, 5),
+    ("ok", None, None),
+])
+def test_model_guide_mismatch(behavior, model_size, guide_size, model):
+    model = poutine.trace(model)
+    expected_ind = model(guide_size)
+    if behavior == "ok":
+        actual_ind = poutine.replay(model, model.trace)(model_size)
+        assert actual_ind == expected_ind
+    else:
+        with pytest.raises(ValueError):
+            poutine.replay(model, model.trace)(model_size)
diff --git a/tests/poutine/test_poutines.py b/tests/poutine/test_poutines.py
--- a/tests/poutine/test_poutines.py
+++ b/tests/poutine/test_poutines.py
@@ -438,10 +438,19 @@ def inner(i, _i, _x):
             return mu_latent
 
         self.model = model
-        self.expected_nodes = set(['z_0_0', 'z_0_1', 'z_1_0', 'z_1_1', 'mu_latent',
+        self.expected_nodes = set(["z_0_0", "z_0_1", "z_1_0", "z_1_1", "mu_latent",
+                                   "map_outer", "map_inner_0", "map_inner_1",
                                    "_INPUT", "_RETURN"])
-        self.expected_edges = set([('mu_latent', 'z_0_0'), ('mu_latent', 'z_0_1'),
-                                   ('mu_latent', 'z_1_0'), ('mu_latent', 'z_1_1')])
+        self.expected_edges = set([
+            ("mu_latent", "z_0_0"), ("mu_latent", "z_0_1"),
+            ("mu_latent", "z_1_0"), ("mu_latent", "z_1_1"),
+            ("map_inner_0", "z_0_0"), ("map_inner_0", "z_0_1"),
+            ("map_inner_1", "z_1_0"), ("map_inner_1", "z_1_1"),
+            ("map_outer", "map_inner_0"), ("map_outer", "map_inner_1"),
+            ("map_outer", "z_0_0"), ("map_outer", "z_0_1"),
+            ("map_outer", "z_1_0"), ("map_outer", "z_1_1"),
+            ("mu_latent", "map_inner_0"), ("mu_latent", "map_inner_1"), ("mu_latent", "map_outer"),
+        ])
 
     def test_graph_structure(self):
         tracegraph = poutine.trace(self.model, graph_type="dense").get_trace()