Repository: pyro-ppl/pyro

Issue:
remove map_data
it's time.


Discussion:
I'm mostly looking forward to updating tests to use `irange` and `iarange` rather than `map_data`.
Yep!

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -281,33 +281,6 @@ def irange(name, size, subsample_size=None, subsample=None, use_cuda=None):
                     yield i.item() if isinstance(i, Variable) else i
 
 
-def map_data(name, data, fn, batch_size=None, batch_dim=0, use_cuda=None):
-    """
-    Data subsampling with the important property that all the data are conditionally independent.
-
-    With default values of `batch_size` and `batch_dim`, `map_data` behaves like `map`.
-    More precisely, `map_data('foo', data, fn)` is equivalent to `[fn(i, x) for i, x in enumerate(data)]`.
-
-    :param str name: named argument
-    :param data: data to subsample
-    :param callable fn: a function taking `(index, datum)` pairs, where `dataum = data[index]`
-    :param int batch_size: number of samples per batch, or zero for the entire dataset
-    :param int batch_dim: dimension to subsample for tensor inputs
-    :param bool use_cuda: Optional bool specifying whether to use cuda tensors
-        for `log_pdf`. Defaults to `torch.Tensor.is_cuda`.
-    :return: a list of values returned by `fn`
-    """
-
-    use_cuda = use_cuda or getattr(data, 'is_cuda', None)
-    if isinstance(data, (torch.Tensor, Variable)):
-        size = data.size(batch_dim)
-        with iarange(name, size, batch_size, use_cuda=use_cuda) as batch:
-            return fn(batch, data.index_select(batch_dim, batch))
-    else:
-        size = len(data)
-        return [fn(i, data[i]) for i in irange(name, size, batch_size, use_cuda=use_cuda)]
-
-
 # XXX this should have the same call signature as torch.Tensor constructors
 def param(name, *args, **kwargs):
     """
diff --git a/pyro/infer/tracegraph_elbo.py b/pyro/infer/tracegraph_elbo.py
--- a/pyro/infer/tracegraph_elbo.py
+++ b/pyro/infer/tracegraph_elbo.py
@@ -35,7 +35,7 @@ def _get_baseline_options(site):
 
 
 def _compute_downstream_costs(model_trace, guide_trace,  #
-                              model_vec_md_nodes, guide_vec_md_nodes,  #
+                              model_iarange_nodes, guide_iarange_nodes,  #
                               non_reparam_nodes):
     # recursively compute downstream cost nodes for all sample sites in model and guide
     # (even though ultimately just need for non-reparameterizable sample sites)
@@ -49,7 +49,7 @@ def _compute_downstream_costs(model_trace, guide_trace,  #
 
     downstream_guide_cost_nodes = {}
     downstream_costs = {}
-    stacks = model_trace.graph["vectorized_map_data_info"]['vec_md_stacks']
+    stacks = model_trace.graph["iarange_info"]['iarange_stacks']
 
     def n_compatible_indices(dest_node, source_node):
         n_compatible = 0
@@ -130,7 +130,7 @@ def _compute_elbo_reparam(model_trace, guide_trace, non_reparam_nodes):
     return torch_data_sum(elbo), surrogate_elbo
 
 
-def _compute_elbo_non_reparam(guide_trace, guide_vec_md_nodes,  #
+def _compute_elbo_non_reparam(guide_trace, guide_iarange_nodes,  #
                               non_reparam_nodes, downstream_costs):
     # construct all the reinforce-like terms.
     # we include only downstream costs to reduce variance
@@ -140,7 +140,7 @@ def _compute_elbo_non_reparam(guide_trace, guide_vec_md_nodes,  #
     baseline_loss = 0.0
     for node in non_reparam_nodes:
         guide_site = guide_trace.nodes[node]
-        log_pdf_key = 'batch_log_pdf' if node in guide_vec_md_nodes else 'log_pdf'
+        log_pdf_key = 'batch_log_pdf' if node in guide_iarange_nodes else 'log_pdf'
         downstream_cost = downstream_costs[node]
         baseline = 0.0
         (nn_baseline, nn_baseline_input, use_decaying_avg_baseline, baseline_beta,
@@ -261,9 +261,9 @@ def loss_and_grads(self, model, guide, *args, **kwargs):
         return loss
 
     def _loss_and_grads_particle(self, weight, model_trace, guide_trace):
-        # get info regarding rao-blackwellization of vectorized map_data
-        guide_vec_md_nodes = guide_trace.graph["vectorized_map_data_info"]['nodes']
-        model_vec_md_nodes = model_trace.graph["vectorized_map_data_info"]['nodes']
+        # get info regarding rao-blackwellization of iarange
+        guide_iarange_nodes = guide_trace.graph["iarange_info"]['nodes']
+        model_iarange_nodes = model_trace.graph["iarange_info"]['nodes']
 
         # have the trace compute all the individual (batch) log pdf terms
         # and score function terms (if present) so that they are available below
@@ -284,9 +284,9 @@ def _loss_and_grads_particle(self, weight, model_trace, guide_trace):
         baseline_loss = 0.0
         if non_reparam_nodes:
             downstream_costs, _ = _compute_downstream_costs(
-                    model_trace, guide_trace,  model_vec_md_nodes, guide_vec_md_nodes, non_reparam_nodes)
+                    model_trace, guide_trace,  model_iarange_nodes, guide_iarange_nodes, non_reparam_nodes)
             surrogate_elbo_term, baseline_loss = _compute_elbo_non_reparam(
-                    guide_trace, guide_vec_md_nodes, non_reparam_nodes, downstream_costs)
+                    guide_trace, guide_iarange_nodes, non_reparam_nodes, downstream_costs)
             surrogate_elbo += surrogate_elbo_term
 
         # collect parameters to train from model and guide
diff --git a/pyro/poutine/trace_poutine.py b/pyro/poutine/trace_poutine.py
--- a/pyro/poutine/trace_poutine.py
+++ b/pyro/poutine/trace_poutine.py
@@ -5,16 +5,16 @@
 from .util import site_is_subsample
 
 
-def get_vectorized_map_data_info(trace):
+def get_iarange_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`.
+    This determines whether the :class:`~pyro.iarange`s 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, 'warnings': set()}
-    vec_md_stacks = set()
+    iarange_info = {'rao-blackwellization-condition': True, 'warnings': set()}
+    iarange_stacks = set()
     stack_dict = {}
 
     for name, node in nodes.items():
@@ -22,40 +22,40 @@ def get_vectorized_map_data_info(trace):
             continue
         if node["type"] in ("sample", "param"):
             stack = node["cond_indep_stack"]
-            vec_mds = [x for x in stack if x.vectorized]
-            stack_dict[name] = vec_mds
+            iarange_stack = [x for x in stack if x.vectorized]
+            stack_dict[name] = iarange_stack
 
     for name, node in nodes.items():
         if site_is_subsample(node):
             continue
         if node["type"] in ("sample", "param"):
             stack = node["cond_indep_stack"]
-            vec_mds = [x for x in stack if x.vectorized]
-            stack_dict[name] = vec_mds
-            # check for nested vectorized map datas
-            if len(vec_mds) > 1:
-                vectorized_map_data_info['rao-blackwellization-condition'] = False
-                vectorized_map_data_info['warnings'].add('nested iarange')
-            # check that vectorized map datas only found at innermost position
-            if vec_mds and not stack[-1].vectorized:
-                vectorized_map_data_info['rao-blackwellization-condition'] = False
-                vectorized_map_data_info['warnings'].add('non-leaf iarange')
-            # enforce that if there are multiple vectorized map_datas, they are all
-            # independent of one another because of enclosing list map_datas
+            iarange_stack = [x for x in stack if x.vectorized]
+            stack_dict[name] = iarange_stack
+            # check for nested iaranges
+            if len(iarange_stack) > 1:
+                iarange_info['rao-blackwellization-condition'] = False
+                iarange_info['warnings'].add('nested iarange')
+            # check that iaranges only found at innermost position
+            if iarange_stack and not stack[-1].vectorized:
+                iarange_info['rao-blackwellization-condition'] = False
+                iarange_info['warnings'].add('non-leaf iarange')
+            # enforce that if there are multiple iaranges, they are all
+            # independent of one another because of enclosing iranges
             # (step 1: collect the stacks)
-            if vec_mds:
-                vec_md_stacks.add(stack)
+            if iarange_stack:
+                iarange_stacks.add(stack)
             # bail, since condition false
-            if not vectorized_map_data_info['rao-blackwellization-condition']:
+            if not iarange_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
+    # enforce that if there are multiple iaranges, they are all
+    # independent of one another because of enclosing iranges
     # (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):
+    if iarange_info['rao-blackwellization-condition']:
+        iarange_stacks = list(iarange_stacks)
+        for i, stack_i in enumerate(iarange_stacks):
+            for j, stack_j in enumerate(iarange_stacks):
                 # only check unique pairs
                 if i <= j:
                     continue
@@ -64,23 +64,23 @@ def get_vectorized_map_data_info(trace):
                     if md_i.name == md_j.name and md_i.counter != md_j.counter:
                         ij_independent = True
                 if not ij_independent:
-                    vectorized_map_data_info['rao-blackwellization-condition'] = False
-                    vectorized_map_data_info['warnings'].add('there exist dependent iaranges')
+                    iarange_info['rao-blackwellization-condition'] = False
+                    iarange_info['warnings'].add('there exist dependent iaranges')
                     break
 
-    vec_md_stacks = list(vec_md_stacks)
-    vectorized_map_data_info['vec_md_stacks'] = stack_dict
+    iarange_stacks = list(iarange_stacks)
+    iarange_info['iarange_stacks'] = stack_dict
 
     # construct data structure consumed by tracegraph_kl_qp
-    vectorized_map_data_info['nodes'] = set()
+    iarange_info['nodes'] = set()
     for name, node in nodes.items():
         if site_is_subsample(node):
             continue
         if node["type"] in ("sample", "param"):
             if any(x.vectorized for x in node["cond_indep_stack"]):
-                vectorized_map_data_info['nodes'].add(name)
+                iarange_info['nodes'].add(name)
 
-    return vectorized_map_data_info
+    return iarange_info
 
 
 def identify_dense_edges(trace):
@@ -137,8 +137,7 @@ def __exit__(self, *args, **kwargs):
         """
         if self.graph_type == "dense":
             identify_dense_edges(self.trace)
-            self.trace.graph["vectorized_map_data_info"] = \
-                get_vectorized_map_data_info(self.trace)
+            self.trace.graph["iarange_info"] = get_iarange_info(self.trace)
         return super(TraceMessenger, self).__exit__(*args, **kwargs)
 
     def get_trace(self):


Test patch:
diff --git a/tests/infer/test_compute_downstream_costs.py b/tests/infer/test_compute_downstream_costs.py
--- a/tests/infer/test_compute_downstream_costs.py
+++ b/tests/infer/test_compute_downstream_costs.py
@@ -17,12 +17,12 @@
 
 
 def _brute_force_compute_downstream_costs(model_trace, guide_trace,  #
-                                          model_vec_md_nodes, guide_vec_md_nodes,  #
+                                          model_iarange_nodes, guide_iarange_nodes,  #
                                           non_reparam_nodes):
 
     guide_nodes = [x for x in guide_trace.nodes if guide_trace.nodes[x]["type"] == "sample"]
     downstream_costs, downstream_guide_cost_nodes = {}, {}
-    stacks = model_trace.graph["vectorized_map_data_info"]['vec_md_stacks']
+    stacks = model_trace.graph["iarange_info"]['iarange_stacks']
 
     def n_compatible_indices(dest_node, source_node):
         n_compatible = 0
@@ -124,21 +124,21 @@ def test_compute_downstream_costs_big_model_guide_pair(include_inner_1, include_
     model_trace.compute_batch_log_pdf()
     guide_trace.compute_batch_log_pdf()
 
-    guide_vec_md_info = guide_trace.graph["vectorized_map_data_info"]
-    model_vec_md_info = model_trace.graph["vectorized_map_data_info"]
-    guide_vec_md_condition = guide_vec_md_info['rao-blackwellization-condition']
-    model_vec_md_condition = model_vec_md_info['rao-blackwellization-condition']
-    do_vec_rb = guide_vec_md_condition and model_vec_md_condition
-    guide_vec_md_nodes = guide_vec_md_info['nodes'] if do_vec_rb else set()
-    model_vec_md_nodes = model_vec_md_info['nodes'] if do_vec_rb else set()
+    guide_iarange_info = guide_trace.graph["iarange_info"]
+    model_iarange_info = model_trace.graph["iarange_info"]
+    guide_iarange_condition = guide_iarange_info['rao-blackwellization-condition']
+    model_iarange_condition = model_iarange_info['rao-blackwellization-condition']
+    do_vec_rb = guide_iarange_condition and model_iarange_condition
+    guide_iarange_nodes = guide_iarange_info['nodes'] if do_vec_rb else set()
+    model_iarange_nodes = model_iarange_info['nodes'] if do_vec_rb else set()
     non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
 
     dc, dc_nodes = _compute_downstream_costs(model_trace, guide_trace,
-                                             model_vec_md_nodes, guide_vec_md_nodes,
+                                             model_iarange_nodes, guide_iarange_nodes,
                                              non_reparam_nodes)
 
     dc_brute, dc_nodes_brute = _brute_force_compute_downstream_costs(model_trace, guide_trace,
-                                                                     model_vec_md_nodes, guide_vec_md_nodes,
+                                                                     model_iarange_nodes, guide_iarange_nodes,
                                                                      non_reparam_nodes)
 
     assert dc_nodes == dc_nodes_brute
@@ -257,21 +257,21 @@ def test_compute_downstream_costs_duplicates(dim):
     model_trace.compute_batch_log_pdf()
     guide_trace.compute_batch_log_pdf()
 
-    guide_vec_md_info = guide_trace.graph["vectorized_map_data_info"]
-    model_vec_md_info = model_trace.graph["vectorized_map_data_info"]
-    guide_vec_md_condition = guide_vec_md_info['rao-blackwellization-condition']
-    model_vec_md_condition = model_vec_md_info['rao-blackwellization-condition']
-    do_vec_rb = guide_vec_md_condition and model_vec_md_condition
-    guide_vec_md_nodes = guide_vec_md_info['nodes'] if do_vec_rb else set()
-    model_vec_md_nodes = model_vec_md_info['nodes'] if do_vec_rb else set()
+    guide_iarange_info = guide_trace.graph["iarange_info"]
+    model_iarange_info = model_trace.graph["iarange_info"]
+    guide_iarange_condition = guide_iarange_info['rao-blackwellization-condition']
+    model_iarange_condition = model_iarange_info['rao-blackwellization-condition']
+    do_vec_rb = guide_iarange_condition and model_iarange_condition
+    guide_iarange_nodes = guide_iarange_info['nodes'] if do_vec_rb else set()
+    model_iarange_nodes = model_iarange_info['nodes'] if do_vec_rb else set()
     non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
 
     dc, dc_nodes = _compute_downstream_costs(model_trace, guide_trace,
-                                             model_vec_md_nodes, guide_vec_md_nodes,
+                                             model_iarange_nodes, guide_iarange_nodes,
                                              non_reparam_nodes)
 
     dc_brute, dc_nodes_brute = _brute_force_compute_downstream_costs(model_trace, guide_trace,
-                                                                     model_vec_md_nodes, guide_vec_md_nodes,
+                                                                     model_iarange_nodes, guide_iarange_nodes,
                                                                      non_reparam_nodes)
 
     assert dc_nodes == dc_nodes_brute
@@ -329,21 +329,21 @@ def test_compute_downstream_costs_iarange_in_irange(dim1):
     model_trace.compute_batch_log_pdf()
     guide_trace.compute_batch_log_pdf()
 
-    guide_vec_md_info = guide_trace.graph["vectorized_map_data_info"]
-    model_vec_md_info = model_trace.graph["vectorized_map_data_info"]
-    guide_vec_md_condition = guide_vec_md_info['rao-blackwellization-condition']
-    model_vec_md_condition = model_vec_md_info['rao-blackwellization-condition']
-    do_vec_rb = guide_vec_md_condition and model_vec_md_condition
-    guide_vec_md_nodes = guide_vec_md_info['nodes'] if do_vec_rb else set()
-    model_vec_md_nodes = model_vec_md_info['nodes'] if do_vec_rb else set()
+    guide_iarange_info = guide_trace.graph["iarange_info"]
+    model_iarange_info = model_trace.graph["iarange_info"]
+    guide_iarange_condition = guide_iarange_info['rao-blackwellization-condition']
+    model_iarange_condition = model_iarange_info['rao-blackwellization-condition']
+    do_vec_rb = guide_iarange_condition and model_iarange_condition
+    guide_iarange_nodes = guide_iarange_info['nodes'] if do_vec_rb else set()
+    model_iarange_nodes = model_iarange_info['nodes'] if do_vec_rb else set()
     non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
 
     dc, dc_nodes = _compute_downstream_costs(model_trace, guide_trace,
-                                             model_vec_md_nodes, guide_vec_md_nodes,
+                                             model_iarange_nodes, guide_iarange_nodes,
                                              non_reparam_nodes)
 
     dc_brute, dc_nodes_brute = _brute_force_compute_downstream_costs(model_trace, guide_trace,
-                                                                     model_vec_md_nodes, guide_vec_md_nodes,
+                                                                     model_iarange_nodes, guide_iarange_nodes,
                                                                      non_reparam_nodes)
 
     assert dc_nodes == dc_nodes_brute
@@ -400,21 +400,21 @@ def test_compute_downstream_costs_irange_in_iarange(dim1, dim2):
     model_trace.compute_batch_log_pdf()
     guide_trace.compute_batch_log_pdf()
 
-    guide_vec_md_info = guide_trace.graph["vectorized_map_data_info"]
-    model_vec_md_info = model_trace.graph["vectorized_map_data_info"]
-    guide_vec_md_condition = guide_vec_md_info['rao-blackwellization-condition']
-    model_vec_md_condition = model_vec_md_info['rao-blackwellization-condition']
-    do_vec_rb = guide_vec_md_condition and model_vec_md_condition
-    guide_vec_md_nodes = guide_vec_md_info['nodes'] if do_vec_rb else set()
-    model_vec_md_nodes = model_vec_md_info['nodes'] if do_vec_rb else set()
+    guide_iarange_info = guide_trace.graph["iarange_info"]
+    model_iarange_info = model_trace.graph["iarange_info"]
+    guide_iarange_condition = guide_iarange_info['rao-blackwellization-condition']
+    model_iarange_condition = model_iarange_info['rao-blackwellization-condition']
+    do_vec_rb = guide_iarange_condition and model_iarange_condition
+    guide_iarange_nodes = guide_iarange_info['nodes'] if do_vec_rb else set()
+    model_iarange_nodes = model_iarange_info['nodes'] if do_vec_rb else set()
     non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
 
     dc, dc_nodes = _compute_downstream_costs(model_trace, guide_trace,
-                                             model_vec_md_nodes, guide_vec_md_nodes,
+                                             model_iarange_nodes, guide_iarange_nodes,
                                              non_reparam_nodes)
 
     dc_brute, dc_nodes_brute = _brute_force_compute_downstream_costs(model_trace, guide_trace,
-                                                                     model_vec_md_nodes, guide_vec_md_nodes,
+                                                                     model_iarange_nodes, guide_iarange_nodes,
                                                                      non_reparam_nodes)
 
     assert dc_nodes == dc_nodes_brute
diff --git a/tests/infer/test_elbo_mapdata.py b/tests/infer/test_elbo_mapdata.py
--- a/tests/infer/test_elbo_mapdata.py
+++ b/tests/infer/test_elbo_mapdata.py
@@ -41,6 +41,7 @@ def add_data_point(x, y):
     add_data_point(0.09, 0.41)
     add_data_point(-0.04, 0.17)
 
+    data = torch.stack(data)
     n_data = Variable(torch.Tensor([len(data)]))
     analytic_lam_n = lam0 + n_data.expand_as(lam) * lam
     analytic_log_sig_n = -0.5 * torch.log(analytic_lam_n)
@@ -55,20 +56,13 @@ def model():
         mu_latent = pyro.sample("mu_latent",
                                 dist.Normal(mu0, torch.pow(lam0, -0.5)).reshape(extra_event_dims=1))
         if map_type == "irange":
-            pyro.map_data("aaa", data,
-                          lambda i, x: pyro.sample("obs_%d" % i,
-                                                   dist.Normal(mu_latent, torch.pow(lam, -0.5))
-                                                       .reshape(extra_event_dims=1),
-                                                   obs=x),
-                          batch_size=batch_size)
+            for i in pyro.irange("aaa", len(data), batch_size):
+                pyro.sample("obs_%d" % i, dist.Normal(mu_latent, torch.pow(lam, -0.5)) .reshape(extra_event_dims=1),
+                            obs=data[i]),
         elif map_type == "iarange":
-            tdata = torch.cat([xi.view(1, -1) for xi in data], 0)
-            pyro.map_data("aaa", tdata,
-                          lambda i, x: pyro.sample("obs",
-                                                   dist.Normal(mu_latent, torch.pow(lam, -0.5))
-                                                       .reshape(extra_event_dims=1),
-                                                   obs=x),
-                          batch_size=batch_size)
+            with pyro.iarange("aaa", len(data), batch_size) as ind:
+                pyro.sample("obs", dist.Normal(mu_latent, torch.pow(lam, -0.5)) .reshape(extra_event_dims=1),
+                            obs=data[ind]),
         else:
             for i, x in enumerate(data):
                 pyro.observe('obs_%d' % i,
@@ -86,11 +80,12 @@ def guide():
         sig_q = torch.exp(log_sig_q)
         pyro.sample("mu_latent", dist.Normal(mu_q, sig_q).reshape(extra_event_dims=1))
         if map_type == "irange" or map_type is None:
-            pyro.map_data("aaa", data, lambda i, x: None, batch_size=batch_size)
+            for i in pyro.irange("aaa", len(data), batch_size):
+                pass
         elif map_type == "iarange":
-            tdata = torch.cat([xi.view(1, -1) for xi in data], 0)
-            # dummy map_data to do subsampling for observe
-            pyro.map_data("aaa", tdata, lambda i, x: None, batch_size=batch_size)
+            # dummy iarange to do subsampling for observe
+            with pyro.iarange("aaa", len(data), batch_size):
+                pass
         else:
             pass
 
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
@@ -432,7 +432,6 @@ def guide():
         assert_equal(0.0, tau_error, prec=0.05)
 
 
-@pytest.mark.xfail(reason="stale; needs to be rewritten using irange, iarange")
 @pytest.mark.init(rng_seed=0)
 @pytest.mark.stage("integration", "integration_batch_1")
 class RaoBlackwellizationTests(TestCase):
@@ -458,8 +457,8 @@ def setUp(self):
         self.analytic_mu_n = self.sum_data * (self.lam / self.analytic_lam_n) +\
             self.mu0 * (self.lam0 / self.analytic_lam_n)
 
-    # this tests rao-blackwellization in elbo for nested list map_datas
-    def test_nested_list_map_data_in_elbo(self, n_steps=4000):
+    # this tests rao-blackwellization in elbo for nested iranges
+    def test_nested_irange_in_elbo(self, n_steps=4000):
         pyro.clear_param_store()
 
         def model():
@@ -467,38 +466,23 @@ def model():
                                     fakes.NonreparameterizedNormal(self.mu0,
                                                                    torch.pow(self.lam0,
                                                                              -0.5)).reshape(extra_event_dims=1))
-
-            def obs_outer(i, x):
-                pyro.map_data("map_obs_inner_%d" % i, x, lambda _i, _x:
-                              obs_inner(i, _i, _x), batch_size=3)
-
-            def obs_inner(i, _i, _x):
-                pyro.sample("obs_%d_%d" % (i, _i),
-                            dist.Normal(mu_latent, torch.pow(self.lam, -0.5)).reshape(extra_event_dims=1),
-                            obs=_x)
-
-            pyro.map_data("map_obs_outer", self.data, lambda i, x:
-                          obs_outer(i, x), batch_size=3)
-
-            return mu_latent
+            for i in pyro.irange("outer", self.n_outer):
+                for j in pyro.irange("inner", self.n_inner):
+                    pyro.sample("obs_%d_%d" % (i, j),
+                                dist.Normal(mu_latent, torch.pow(self.lam, -0.5)).reshape(extra_event_dims=1),
+                                obs=self.data[i][j])
 
         def guide():
             mu_q = pyro.param("mu_q", variable(self.analytic_mu_n.expand(2) + 0.234, requires_grad=True))
             log_sig_q = pyro.param("log_sig_q",
                                    variable(self.analytic_log_sig_n.expand(2) - 0.27, requires_grad=True))
             sig_q = torch.exp(log_sig_q)
-            mu_latent = pyro.sample("mu_latent", fakes.NonreparameterizedNormal(mu_q,
-                                                                                sig_q).reshape(extra_event_dims=1),
-                                    infer=dict(baseline=dict(use_decaying_avg_baseline=True)))
-
-            def obs_outer(i, x):
-                pyro.map_data("map_obs_inner_%d" % i, x, lambda _i, _x:
-                              None, batch_size=3)
-
-            pyro.map_data("map_obs_outer", self.data, lambda i, x:
-                          obs_outer(i, x), batch_size=3)
+            pyro.sample("mu_latent", fakes.NonreparameterizedNormal(mu_q, sig_q).reshape(extra_event_dims=1),
+                        infer=dict(baseline=dict(use_decaying_avg_baseline=True)))
 
-            return mu_latent
+            for i in pyro.irange("outer", self.n_outer):
+                for j in pyro.irange("inner", self.n_inner):
+                    pass
 
         guide_trace = pyro.poutine.trace(guide, graph_type="dense").get_trace()
         model_trace = pyro.poutine.trace(pyro.poutine.replay(model, guide_trace),
@@ -521,13 +505,13 @@ def obs_outer(i, x):
         assert_equal(0.0, mu_error, prec=0.04)
         assert_equal(0.0, log_sig_error, prec=0.04)
 
-    # this tests rao-blackwellization and baselines for a vectorized map_data
-    # inside of a list map_data with superfluous random variables to complexify the
+    # this tests rao-blackwellization and baselines for iarange
+    # inside of an irange with superfluous random variables to complexify the
     # graph structure and introduce additional baselines
-    def test_vectorized_map_data_in_elbo_with_superfluous_rvs(self):
-        self._test_vectorized_map_data_in_elbo(n_superfluous_top=1, n_superfluous_bottom=1, n_steps=5000)
+    def test_iarange_in_elbo_with_superfluous_rvs(self):
+        self._test_iarange_in_elbo(n_superfluous_top=1, n_superfluous_bottom=1, n_steps=5000)
 
-    def _test_vectorized_map_data_in_elbo(self, n_superfluous_top, n_superfluous_bottom, n_steps):
+    def _test_iarange_in_elbo(self, n_superfluous_top, n_superfluous_bottom, n_steps):
         pyro.clear_param_store()
         self.data_tensor = Variable(torch.zeros(9, 2))
         for _out in range(self.n_outer):
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
@@ -14,41 +14,16 @@
 logger = logging.getLogger(__name__)
 
 
-@pytest.mark.parametrize("batch_dim", [0, 1])
-def test_batch_dim(batch_dim):
-
-    data = Variable(torch.randn(4, 5, 7))
-
-    def local_model(ixs, _xs):
-        xs = _xs.view(-1, _xs.size(2))
-        return pyro.sample("xs", dist.Normal(xs, Variable(torch.ones(xs.size()))))
-
-    def model():
-        return pyro.map_data("md", data, local_model,
-                             batch_size=1, batch_dim=batch_dim)
-
-    tr = poutine.trace(model).get_trace()
-    assert tr.nodes["xs"]["value"].size(0) == data.size(1 - batch_dim)
-    assert tr.nodes["xs"]["value"].size(1) == data.size(2)
-
-
-def test_nested_map_data():
+def test_nested_irange():
     means = [Variable(torch.randn(2)) for i in range(8)]
     mean_batch_size = 2
     stds = [Variable(torch.abs(torch.randn(2))) for i in range(6)]
     std_batch_size = 3
 
     def model(means, stds):
-        return pyro.map_data("a", means,
-                             lambda i, x:
-                             pyro.map_data("a_{}".format(i), stds,
-                                           lambda j, y:
-                                           pyro.sample("x_{}{}".format(i, j),
-                                                       dist.Normal(x, y)),
-                                           batch_size=std_batch_size),
-                             batch_size=mean_batch_size)
-
-    model = model
+        return [[pyro.sample("x_{}{}".format(i, j), dist.Normal(means[i], stds[j]))
+                 for j in pyro.irange("b", len(stds), std_batch_size)]
+                for i in pyro.irange("a", len(means), mean_batch_size)]
 
     xs = model(means, stds)
     assert len(xs) == mean_batch_size
@@ -91,39 +66,9 @@ def nested_irange_model(subsample_size):
     return result
 
 
-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.Normal(mu[batch], sigma[batch]))
-        return batch
-
-    LongTensor = torch.cuda.LongTensor if torch.Tensor.is_cuda else torch.LongTensor
-    ind = Variable(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.Normal(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,
-    nested_irange_model,
-    map_data_vector_model,
-    map_data_iter_model,
-], ids=['iarange', 'irange', 'nested_irange', 'map_data_vector', 'map_data_iter'])
+@pytest.mark.parametrize('model', [iarange_model, irange_model, nested_irange_model],
+                         ids=['iarange', 'irange', 'nested_irange'])
 def test_cond_indep_stack(model, subsample_size):
     tr = poutine.trace(model).get_trace(subsample_size)
     for name, node in tr.nodes.items():
@@ -132,13 +77,8 @@ def test_cond_indep_stack(model, subsample_size):
 
 
 @pytest.mark.parametrize('subsample_size', [5, 20])
-@pytest.mark.parametrize('model', [
-    iarange_model,
-    irange_model,
-    nested_irange_model,
-    map_data_vector_model,
-    map_data_iter_model,
-], ids=['iarange', 'irange', 'nested_irange', 'map_data_vector', 'map_data_iter'])
+@pytest.mark.parametrize('model', [iarange_model, irange_model, nested_irange_model],
+                         ids=['iarange', 'irange', 'nested_irange'])
 def test_replay(model, subsample_size):
     pyro.set_rng_seed(0)
 
@@ -190,41 +130,15 @@ def irange_cuda_model(subsample_size):
         pyro.sample("x_{}".format(i), dist.Normal(mu[i], sigma[i]))
 
 
-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.Normal(mu, sigma[i])),
-                  batch_size=subsample_size, use_cuda=True)
-
-
-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.Normal(mu, sigma[i])),
-                  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"])
+@pytest.mark.parametrize('model', [iarange_cuda_model, irange_cuda_model], ids=["iarange", "irange"])
 def test_cuda(model, subsample_size):
     tr = poutine.trace(model).get_trace(subsample_size)
     assert tr.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('model', [iarange_model, irange_model], ids=['iarange', 'irange'])
 @pytest.mark.parametrize("behavior,model_size,guide_size", [
     ("error", 20, 5),
     ("error", 5, 20),
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
@@ -433,19 +433,11 @@ class IndirectLambdaPoutineTests(TestCase):
     def setUp(self):
 
         def model(batch_size_outer=2, batch_size_inner=2):
+            data = [[ng_ones(1)] * 2] * 2
             mu_latent = pyro.sample("mu_latent", dist.Normal(ng_zeros(1), ng_ones(1)))
-
-            def outer(i, x):
-                pyro.map_data("map_inner_%d" % i, x, lambda _i, _x:
-                              inner(i, _i, _x), batch_size=batch_size_inner)
-
-            def inner(i, _i, _x):
-                pyro.sample("z_%d_%d" % (i, _i), dist.Normal(mu_latent + _x, ng_ones(1)))
-
-            pyro.map_data("map_outer", [[ng_ones(1)] * 2] * 2, lambda i, x:
-                          outer(i, x), batch_size=batch_size_outer)
-
-            return mu_latent
+            for i in pyro.irange("irange_outer", 2, batch_size_outer):
+                for j in pyro.irange("irange_inner_%d" % i, 2, batch_size_inner):
+                    pyro.sample("z_%d_%d" % (i, j), dist.Normal(mu_latent + data[i][j], ng_ones(1)))
 
         self.model = model
         self.expected_nodes = set(["z_0_0", "z_0_1", "z_1_0", "z_1_1", "mu_latent",
@@ -457,10 +449,10 @@ def inner(i, _i, _x):
 
     def test_graph_structure(self):
         tracegraph = poutine.trace(self.model, graph_type="dense").get_trace()
-        # Ignore structure on map_* nodes.
-        actual_nodes = set(n for n in tracegraph.nodes() if not n.startswith("map_"))
+        # Ignore structure on irange_* nodes.
+        actual_nodes = set(n for n in tracegraph.nodes() if not n.startswith("irange_"))
         actual_edges = set((n1, n2) for n1, n2 in tracegraph.edges
-                           if not n1.startswith("map_") if not n2.startswith("map_"))
+                           if not n1.startswith("irange_") if not n2.startswith("irange_"))
         assert actual_nodes == self.expected_nodes
         assert actual_edges == self.expected_edges