Repository: pyro-ppl/pyro

Issue:
more control over subsampling in map_data
currently, mini-batch indices for map_data subsampling are chosen by pyro. they are also chosen in a stateless way with the consequence that there is non-zero probability that certain data points will never be seen even no matter how many epochs you do. this is bad, especially for large data sets. the simplest solution to this problem would be to give the user the ability to specify `batch_indices` directly in `map_data` (a solution that would presumably be psychologically comforting to users who are used to doing this anyway). another solution, which could coexist with the first solution, would be for the backend to implement more subsampling strategies besides the 'naive' one that we currently have. 


Discussion:
If the user wants to directly specify batch_indices, they can just not use `map_data` and write what they want directly.  Maybe we can add an option to the default `Poutine` to sample without replacement in `map_data`, though.

We could also expose the randomness in `map_data` as a Pyro variable, to enable fancy active subsampling strategies.
unless i'm missing something in that case scale factors won't be propagated and any other benefits from map_data (like independency information used for rao-blackwellization) fall by the wayside
Hmm... I guess not, although this seems like a potential source of nasty statistical bugs, especially in the nested case, or when there's a sequence of variables following a `map_data` that depend on all the variables inside.
yeah, that's certainly true. but my general design philosophy remains:

- try to make the default/highlighted interfaces relatively robust to user error
- but in those cases where finer control can be of a lot of practical importance, allow the advanced user finer knobs even if this opens the door to bugs and math errors. potential pitfalls can also be highlighted with friendly warnings as in #164 
i agree with @martinjankowiak. it's pretty common, eg, to want to iterate over the batches of data.

we should also consider doing this extension in such a way that it plays nice with dynamic data loaders for when data is too big to fit in memory....

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -146,7 +146,7 @@ def batch_log_pdf(self, x):
 
 
 @contextlib.contextmanager
-def iarange(name, size, subsample_size=0):
+def iarange(name, size, subsample_size=0, subsample=None):
     """
     Context manager for ranges indexing iid variables, optionally subsampling.
 
@@ -160,6 +160,9 @@ def iarange(name, size, subsample_size=0):
     :param str name: A name that will be used for this site in a Trace.
     :param int size: The size of the collection being subsampled (like `stop` in builtin `range`).
     :param int subsample_size: Size of minibatches used in subsampling. Defaults to `size` if set to 0.
+    :param subsample: Optional custom subsample for user-defined subsampling schemes.
+        If specified, then `subsample_size` will be set to `len(subsample)`.
+    :type subsample: Anything supporting `len()`.
     :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
 
     Examples::
@@ -167,20 +170,30 @@ def iarange(name, size, subsample_size=0):
         # This version is vectorized:
         >>> with iarange('data', 100, subsample_size=10) as batch:
                 observe('obs', normal, data.index_select(0, batch), mu, sigma)
+
         # This version manually iterates through the batch to deal with control flow.
         >>> with iarange('data', 100, subsample_size=10) as batch:
                 for i in batch:
                     if z[i]:  # Prevents vectorization.
                         observe('obs_{}'.format(i), normal, data[i], mu, sigma)
+
+        # This wraps a user-defined subsampling method for use in pyro:
+        >>> with iarange('data', 100, subsample=my_custom_subsample) as batch:
+                assert batch is my_custom_subsample
+                observe('obs', normal, data.index_select(0, batch), mu, sigma)
     """
-    if subsample_size == 0 or subsample_size >= size:
+    if subsample is not None:
+        subsample_size = len(subsample)
+        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.
         yield Variable(torch.LongTensor(list(range(size))))
         return
 
-    subsample = sample(name, _Subsample(size, subsample_size))
+    if subsample is None:
+        subsample = sample(name, _Subsample(size, subsample_size))
     if len(_PYRO_STACK) == 0:
         yield subsample
     else:
@@ -190,9 +203,17 @@ def iarange(name, size, subsample_size=0):
             yield subsample
 
 
-def irange(name, size, subsample_size=0):
+def irange(name, size, subsample_size=0, subsample=None):
     """
-    Non-vectorized version of iarange.
+    Non-vectorized version of `iarange`. See `iarange` for details.
+
+    :param str name: A name that will be used for this site in a Trace.
+    :param int size: The size of the collection being subsampled (like `stop` in builtin `range`).
+    :param int subsample_size: Size of minibatches used in subsampling. Defaults to `size` if set to 0.
+    :param subsample: Optional custom subsample for user-defined subsampling schemes.
+        If specified, then `subsample_size` will be set to `len(subsample)`.
+    :type subsample: Anything supporting `len()`.
+    :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
 
     Examples::
 
@@ -200,10 +221,15 @@ def irange(name, size, subsample_size=0):
                 if z[i]:  # Prevents vectorization.
                     observe('obs_{}'.format(i), normal, data[i], mu, sigma)
     """
-    with iarange(name, size, subsample_size) as batch:
+    if subsample is not None:
+        subsample_size = len(subsample)
+
+    with iarange(name, size, subsample_size, subsample) as batch:
         # Wrap computation in an independence context.
         indep_context = LambdaPoutine(None, name, 1.0, 'list', 0, subsample_size)
-        for i in batch.data:
+        if isinstance(batch, Variable):
+            batch = batch.data
+        for i in batch:
             with indep_context:
                 yield i
 


Test patch:
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
@@ -160,27 +160,59 @@ def model(means, stds):
             assert tr.nodes[name]["scale"] == 4.0 * 2.0
 
 
-def test_replay_iarange():
-    pyro.set_rng_seed(0)
+def iarange_model():
+    with pyro.iarange('iarange', 20, 5) as batch:
+        result = list(batch.data)
+    return result
 
-    def model():
-        with pyro.iarange('iarange', 10, 2) as batch:
-            return list(batch.data)
 
-    traced_model = poutine.trace(model)
-    expected = traced_model()
-    actual = poutine.replay(model, traced_model.trace)()
-    assert actual == expected
+def irange_model():
+    result = []
+    for i in pyro.irange('irange', 20, 5):
+        result.append(i)
+    return result
 
 
-def test_replay_map_data():
-    pyro.set_rng_seed(0)
-    data = list(range(10))
+def map_data_model():
+    return pyro.map_data('mapdata', range(20), lambda i, x: i, batch_size=5)
 
-    def model():
-        return pyro.map_data('mapdata', data, lambda i, x: i, batch_size=2)
+
+@pytest.mark.parametrize('model', [
+    iarange_model,
+    irange_model,
+    map_data_model,
+], ids=['iarange', 'irange', 'map_data'])
+def test_replay(model):
+    pyro.set_rng_seed(0)
 
     traced_model = poutine.trace(model)
-    expected = traced_model()
-    actual = poutine.replay(model, traced_model.trace)()
-    assert actual == expected
+    original = traced_model()
+
+    replayed = poutine.replay(model, traced_model.trace)()
+    assert replayed == original
+
+    different = traced_model()
+    assert different != original
+
+
+def iarange_custom_model(subsample):
+    with pyro.iarange('iarange', 20, subsample=subsample) as batch:
+        result = batch
+    return result
+
+
+def irange_custom_model(subsample):
+    result = []
+    for i in pyro.irange('irange', 20, subsample=subsample):
+        result.append(i)
+    return result
+
+
+@pytest.mark.parametrize('model', [iarange_custom_model, irange_custom_model],
+                         ids=['iarange', 'irange'])
+def test_custom_subsample(model):
+    pyro.set_rng_seed(0)
+
+    subsample = [1, 3, 5, 7]
+    assert model(subsample) == subsample
+    assert poutine.trace(model)(subsample) == subsample