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`.

Changed files:
pyro/__init__.py, tests/poutine/test_mapdata.py, pyro/poutine/trace_poutine.py, pyro/util.py

diff --git a/pyro/__init__.py b/pyro/__init__.py
index cd2ffabe..2628c467 100644
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -204,14 +204,15 @@ def iarange(name, size, subsample_size=0, subsample=None, use_cuda=False):
         assert subsample_size <= size, 'subsample is larger than size'
     elif subsample_size == 0 or subsample_size >= size:
         subsample_size = size
-    if subsample_size == size:
-        # If not subsampling, there is no need to scale and we can ignore the _PYRO_STACK.
-        result = Variable(torch.LongTensor(list(range(size))))
-        yield result.cuda() if use_cuda else result
-        return
 
     if subsample is None:
-        subsample = sample(name, _Subsample(size, subsample_size, use_cuda))
+        if subsample_size == size:
+            subsample = Variable(torch.LongTensor(list(range(size))))
+            if use_cuda:
+                subsample = subsample.cuda()
+        else:
+            subsample = sample(name, _Subsample(size, subsample_size, use_cuda))
+
     if len(_PYRO_STACK) == 0:
         yield subsample
     else:
diff --git a/pyro/poutine/trace_poutine.py b/pyro/poutine/trace_poutine.py
index 39fb4de8..187ed459 100644
--- a/pyro/poutine/trace_poutine.py
+++ b/pyro/poutine/trace_poutine.py
@@ -1,8 +1,102 @@
-from pyro import util
+from collections import defaultdict
+
 from .poutine import Poutine
 from .trace import Trace
 
 
+def get_vectorized_map_data_info(trace):
+    """
+    This determines whether the vectorized map_datas are rao-blackwellizable by
+    `TraceGraph_ELBO`. This also gathers information to be consumed by
+    downstream by `TraceGraph_ELBO`.
+    """
+    nodes = trace.nodes
+
+    vectorized_map_data_info = {'rao-blackwellization-condition': True}
+    vec_md_stacks = set()
+
+    for name, node in nodes.items():
+        if node["type"] in ("sample", "param"):
+            stack = tuple(reversed(node["map_data_stack"]))
+            vec_mds = list(filter(lambda x: x[2] == 'tensor', stack))
+            # check for nested vectorized map datas
+            if len(vec_mds) > 1:
+                vectorized_map_data_info['rao-blackwellization-condition'] = False
+            # check that vectorized map datas only found at innermost position
+            if len(vec_mds) > 0 and stack[-1][2] == 'list':
+                vectorized_map_data_info['rao-blackwellization-condition'] = False
+            # for now enforce batch_dim = 0 for vectorized map_data
+            # since needed batch_log_pdf infrastructure missing
+            if len(vec_mds) > 0 and vec_mds[0][3] != 0:
+                vectorized_map_data_info['rao-blackwellization-condition'] = False
+            # enforce that if there are multiple vectorized map_datas, they are all
+            # independent of one another because of enclosing list map_datas
+            # (step 1: collect the stacks)
+            if len(vec_mds) > 0:
+                vec_md_stacks.add(stack)
+            # bail, since condition false
+            if not vectorized_map_data_info['rao-blackwellization-condition']:
+                break
+
+    # enforce that if there are multiple vectorized map_datas, they are all
+    # independent of one another because of enclosing list map_datas
+    # (step 2: explicitly check this)
+    if vectorized_map_data_info['rao-blackwellization-condition']:
+        vec_md_stacks = list(vec_md_stacks)
+        for i, stack_i in enumerate(vec_md_stacks):
+            for j, stack_j in enumerate(vec_md_stacks):
+                # only check unique pairs
+                if i <= j:
+                    continue
+                ij_independent = False
+                for md_i, md_j in zip(stack_i, stack_j):
+                    if md_i[0] == md_j[0] and md_i[1] != md_j[1]:
+                        ij_independent = True
+                if not ij_independent:
+                    vectorized_map_data_info['rao-blackwellization-condition'] = False
+                    break
+
+    # construct data structure consumed by tracegraph_kl_qp
+    if vectorized_map_data_info['rao-blackwellization-condition']:
+        vectorized_map_data_info['nodes'] = defaultdict(list)
+        for name, node in nodes.items():
+            if node["type"] in ("sample", "param"):
+                stack = tuple(reversed(node["map_data_stack"]))
+                vec_mds = list(filter(lambda x: x[2] == 'tensor', stack))
+                if len(vec_mds) > 0:
+                    node_batch_dim_pair = (name, vec_mds[0][3])
+                    vectorized_map_data_info['nodes'][vec_mds[0][0]].append(node_batch_dim_pair)
+
+    return vectorized_map_data_info
+
+
+def identify_dense_edges(trace):
+    """
+    Method to add all edges based on the map_data_stack information
+    stored at each site.
+    """
+    for name, node in trace.nodes.items():
+        if node["type"] == "sample":
+            # XXX why tuple?
+            map_data_stack = tuple(reversed(node["map_data_stack"]))
+            for past_name, past_node in trace.nodes.items():
+                if past_node["type"] == "sample":
+                    if past_name == name:
+                        break
+                    past_node_independent = False
+                    past_node_map_data_stack = tuple(
+                        reversed(past_node["map_data_stack"]))
+                    for query, target in zip(map_data_stack,
+                                             past_node_map_data_stack):
+                        if query[0] == target[0] and query[1] != target[1]:
+                            past_node_independent = True
+                            break
+                    if not past_node_independent:
+                        trace.add_edge(past_name, name)
+
+    return trace
+
+
 class TracePoutine(Poutine):
     """
     Execution trace poutine.
@@ -35,9 +129,9 @@ class TracePoutine(Poutine):
         upon exiting the context.
         """
         if self.graph_type == "dense":
-            self.trace = util.identify_dense_edges(self.trace)
+            self.trace = identify_dense_edges(self.trace)
             self.trace.graph["vectorized_map_data_info"] = \\
-                util.get_vectorized_map_data_info(self.trace)
+                get_vectorized_map_data_info(self.trace)
         return super(TracePoutine, self).__exit__(*args, **kwargs)
 
     def __call__(self, *args, **kwargs):
diff --git a/pyro/util.py b/pyro/util.py
index 901a4081..1306360b 100644
--- a/pyro/util.py
+++ b/pyro/util.py
@@ -5,7 +5,6 @@ import functools
 import torch
 from torch.autograd import Variable
 from torch.nn import Parameter
-from collections import defaultdict
 
 
 def detach_iterable(iterable):
@@ -316,73 +315,6 @@ def all_escape(trace, msg):
         (msg["name"] not in trace)
 
 
-def get_vectorized_map_data_info(trace):
-    """
-    this determines whether the vectorized map_datas
-    are rao-blackwellizable by tracegraph_kl_qp
-    also gathers information to be consumed by downstream by tracegraph_kl_qp
-    XXX this logic should probably sit elsewhere
-    """
-    nodes = trace.nodes
-
-    vectorized_map_data_info = {'rao-blackwellization-condition': True}
-    vec_md_stacks = set()
-
-    for name, node in nodes.items():
-        if node["type"] in ("sample", "param"):
-            stack = tuple(reversed(node["map_data_stack"]))
-            vec_mds = list(filter(lambda x: x[2] == 'tensor', stack))
-            # check for nested vectorized map datas
-            if len(vec_mds) > 1:
-                vectorized_map_data_info['rao-blackwellization-condition'] = False
-            # check that vectorized map datas only found at innermost position
-            if len(vec_mds) > 0 and stack[-1][2] == 'list':
-                vectorized_map_data_info['rao-blackwellization-condition'] = False
-            # for now enforce batch_dim = 0 for vectorized map_data
-            # since needed batch_log_pdf infrastructure missing
-            if len(vec_mds) > 0 and vec_mds[0][3] != 0:
-                vectorized_map_data_info['rao-blackwellization-condition'] = False
-            # enforce that if there are multiple vectorized map_datas, they are all
-            # independent of one another because of enclosing list map_datas
-            # (step 1: collect the stacks)
-            if len(vec_mds) > 0:
-                vec_md_stacks.add(stack)
-            # bail, since condition false
-            if not vectorized_map_data_info['rao-blackwellization-condition']:
-                break
-
-    # enforce that if there are multiple vectorized map_datas, they are all
-    # independent of one another because of enclosing list map_datas
-    # (step 2: explicitly check this)
-    if vectorized_map_data_info['rao-blackwellization-condition']:
-        vec_md_stacks = list(vec_md_stacks)
-        for i, stack_i in enumerate(vec_md_stacks):
-            for j, stack_j in enumerate(vec_md_stacks):
-                # only check unique pairs
-                if i <= j:
-                    continue
-                ij_independent = False
-                for md_i, md_j in zip(stack_i, stack_j):
-                    if md_i[0] == md_j[0] and md_i[1] != md_j[1]:
-                        ij_independent = True
-                if not ij_independent:
-                    vectorized_map_data_info['rao-blackwellization-condition'] = False
-                    break
-
-    # construct data structure consumed by tracegraph_kl_qp
-    if vectorized_map_data_info['rao-blackwellization-condition']:
-        vectorized_map_data_info['nodes'] = defaultdict(list)
-        for name, node in nodes.items():
-            if node["type"] in ("sample", "param"):
-                stack = tuple(reversed(node["map_data_stack"]))
-                vec_mds = list(filter(lambda x: x[2] == 'tensor', stack))
-                if len(vec_mds) > 0:
-                    node_batch_dim_pair = (name, vec_mds[0][3])
-                    vectorized_map_data_info['nodes'][vec_mds[0][0]].append(node_batch_dim_pair)
-
-    return vectorized_map_data_info
-
-
 def save_visualization(trace, graph_output):
     """
     render graph and save to file
@@ -410,33 +342,6 @@ def save_visualization(trace, graph_output):
     g.render(graph_output, view=False, cleanup=True)
 
 
-def identify_dense_edges(trace):
-    """
-    Method to add all edges based on the map_data_stack information
-    stored at each site.
-    """
-    for name, node in trace.nodes.items():
-        if node["type"] == "sample":
-            # XXX why tuple?
-            map_data_stack = tuple(reversed(node["map_data_stack"]))
-            for past_name, past_node in trace.nodes.items():
-                if past_node["type"] == "sample":
-                    if past_name == name:
-                        break
-                    past_node_independent = False
-                    past_node_map_data_stack = tuple(
-                        reversed(past_node["map_data_stack"]))
-                    for query, target in zip(map_data_stack,
-                                             past_node_map_data_stack):
-                        if query[0] == target[0] and query[1] != target[1]:
-                            past_node_independent = True
-                            break
-                    if not past_node_independent:
-                        trace.add_edge(past_name, name)
-
-    return trace
-
-
 def deep_getattr(obj, name):
     """
     Python getattr() for arbitrarily deep attributes
diff --git a/tests/poutine/test_mapdata.py b/tests/poutine/test_mapdata.py
index b12316cd..1fd11b6c 100644
--- a/tests/poutine/test_mapdata.py
+++ b/tests/poutine/test_mapdata.py
@@ -160,39 +160,82 @@ def test_nested_map_data():
             assert tr.nodes[name]["scale"] == 4.0 * 2.0
 
 
-def iarange_model():
-    with pyro.iarange('iarange', 20, 5) as batch:
+def iarange_model(subsample_size):
+    mu = Variable(torch.zeros(20))
+    sigma = Variable(torch.ones(20))
+    with pyro.iarange('iarange', 20, subsample_size) as batch:
+        pyro.sample("x", dist.diagnormal, mu[batch], sigma[batch])
         result = list(batch.data)
     return result
 
 
-def irange_model():
+def irange_model(subsample_size):
+    mu = Variable(torch.zeros(20))
+    sigma = Variable(torch.ones(20))
     result = []
-    for i in pyro.irange('irange', 20, 5):
+    for i in pyro.irange('irange', 20, subsample_size):
+        pyro.sample("x_{}".format(i), dist.diagnormal, mu[i], sigma[i])
         result.append(i)
     return result
 
 
-def map_data_model():
-    return pyro.map_data('mapdata', range(20), lambda i, x: i, batch_size=5)
+def map_data_vector_model(subsample_size):
+    mu = Variable(torch.zeros(20))
+    sigma = Variable(torch.ones(20))
+
+    def local_model(batch, unused):
+        pyro.sample("x", dist.diagnormal, mu[batch], sigma[batch])
+        return batch
+
+    ind = Variable(torch.LongTensor(range(20)))
+    batch = pyro.map_data('mapdata', ind, local_model, batch_size=subsample_size)
+    return list(batch.data)
+
+
+def map_data_iter_model(subsample_size):
+    mu = Variable(torch.zeros(20))
+    sigma = Variable(torch.ones(20))
+
+    def local_model(i, unused):
+        pyro.sample("x_{}".format(i), dist.diagnormal, mu[i], sigma[i])
+        return i
+
+    return pyro.map_data('mapdata', range(20), local_model, batch_size=subsample_size)
+
+
+@pytest.mark.parametrize('subsample_size', [5, 20])
+@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'])
+def test_map_data_stack(model, subsample_size):
+    tr = poutine.trace(model).get_trace(subsample_size)
+    for name, node in tr.nodes.items():
+        if name.startswith("x"):
+            assert node["map_data_stack"], "missing map_data_stack at node {}".format(name)
 
 
+@pytest.mark.parametrize('subsample_size', [5, 20])
 @pytest.mark.parametrize('model', [
     iarange_model,
     irange_model,
-    map_data_model,
-], ids=['iarange', 'irange', 'map_data'])
-def test_replay(model):
+    map_data_vector_model,
+    map_data_iter_model,
+], ids=['iarange', 'irange', 'map_data_vector', 'map_data_iter'])
+def test_replay(model, subsample_size):
     pyro.set_rng_seed(0)
 
     traced_model = poutine.trace(model)
-    original = traced_model()
+    original = traced_model(subsample_size)
 
-    replayed = poutine.replay(model, traced_model.trace)()
+    replayed = poutine.replay(model, traced_model.trace)(subsample_size)
     assert replayed == original
 
-    different = traced_model()
-    assert different != original
+    if subsample_size < 20:
+        different = traced_model(subsample_size)
+        assert different != original
 
 
 def iarange_custom_model(subsample):
@@ -218,44 +261,45 @@ def test_custom_subsample(model):
     assert poutine.trace(model)(subsample) == subsample
 
 
-def iarange_cuda_model():
+def iarange_cuda_model(subsample_size):
     mu = Variable(torch.zeros(20).cuda())
     sigma = Variable(torch.ones(20).cuda())
-    with pyro.iarange("data", 20, 5, use_cuda=True) as batch:
+    with pyro.iarange("data", 20, subsample_size, use_cuda=True) as batch:
         pyro.sample("x", dist.diagnormal, mu[batch], sigma[batch])
 
 
-def irange_cuda_model():
+def irange_cuda_model(subsample_size):
     mu = Variable(torch.zeros(20).cuda())
     sigma = Variable(torch.ones(20).cuda())
-    for i in pyro.irange("data", 20, 5, use_cuda=True):
+    for i in pyro.irange("data", 20, subsample_size, use_cuda=True):
         pyro.sample("x_{}".format(i), dist.diagnormal, mu[i], sigma[i])
 
 
-def map_data_vector_cuda_model():
+def map_data_vector_cuda_model(subsample_size):
     mu = Variable(torch.zeros(20).cuda())
     sigma = Variable(torch.ones(20).cuda())
     pyro.map_data("data", mu,
                   lambda i, mu: pyro.sample("x", dist.diagnormal, mu, sigma[i]),
-                  use_cuda=True)
+                  batch_size=subsample_size, use_cuda=True)
 
 
-def map_data_iter_cuda_model():
+def map_data_iter_cuda_model(subsample_size):
     mu = Variable(torch.zeros(20).cuda())
     sigma = Variable(torch.ones(20).cuda())
     pyro.map_data("data", list(mu),
                   lambda i, mu: pyro.sample("x_{}".format(i), dist.diagnormal, mu, sigma[i]),
-                  use_cuda=True)
+                  batch_size=subsample_size, use_cuda=True)
 
 
 @requires_cuda
+@pytest.mark.parametrize('subsample_size', [5, 20])
 @pytest.mark.parametrize('model', [
     iarange_cuda_model,
     irange_cuda_model,
     map_data_vector_cuda_model,
     map_data_iter_cuda_model,
 ], ids=["iarange", "irange", "map_data_vector", "map_data_iter"])
-def test_cuda(model):
-    tr = poutine.trace(model).get_trace()
+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