Repository: pyro-ppl/pyro

Issue:
Batch indices are no longer shared between model and guide
Since 97eac68, the model and guide sub sample data independently, resulting in the model and the guide seeing different mini batches within a single execution. Unless I'm missing something, this looks like a bug.

e.g.

```python
import pyro
import torch
from torch.autograd import Variable
import torch.optim as optim
from pyro.infer.kl_qp import KL_QP

torch.manual_seed(0)
data = range(10)
def model():
    pyro.map_data('mapdata', data, lambda i, x: print('i={},x={}'.format(i, x)), batch_size=2)
    print('-----')
KL_QP(model, model, pyro.optim(optim.Adam, {})).step()
```

At a407f13:

```
i=4,x=4
i=1,x=1
-----
i=4,x=4
i=1,x=1
-----
# cruft ...
```

At 97eac68:

```
i=4,x=4
i=1,x=1
-----
i=3,x=3
i=9,x=9
-----
# cruft ...
```

Batch indices are no longer shared between model and guide
Since 97eac68, the model and guide sub sample data independently, resulting in the model and the guide seeing different mini batches within a single execution. Unless I'm missing something, this looks like a bug.

e.g.

```python
import pyro
import torch
from torch.autograd import Variable
import torch.optim as optim
from pyro.infer.kl_qp import KL_QP

torch.manual_seed(0)
data = range(10)
def model():
    pyro.map_data('mapdata', data, lambda i, x: print('i={},x={}'.format(i, x)), batch_size=2)
    print('-----')
KL_QP(model, model, pyro.optim(optim.Adam, {})).step()
```

At a407f13:

```
i=4,x=4
i=1,x=1
-----
i=4,x=4
i=1,x=1
-----
# cruft ...
```

At 97eac68:

```
i=4,x=4
i=1,x=1
-----
i=3,x=3
i=9,x=9
-----
# cruft ...
```



Discussion:
Uh oh, this is important. Good catch!

@fritzo wanna patch and add a test that would've caught it?
Yes, I'll look into it. Most likely culprit is #184 
Ok sorry, I accidentally removed replay information in #184. I'll add it back. Here's a minimal test case:
```py
def test_replay_map_data():
    pyro.set_rng_seed(0)
    data = list(range(10))

    def model():
        return pyro.map_data('mapdata', data, lambda i, x: i, batch_size=2)

    traced_model = poutine.trace(model)
    expected = traced_model()
    actual = poutine.replay(model, traced_model.trace)()
    assert actual == expected
```
@null-a Thanks for providing such an easily reproducible error!
Uh oh, this is important. Good catch!

@fritzo wanna patch and add a test that would've caught it?
Yes, I'll look into it. Most likely culprit is #184 
Ok sorry, I accidentally removed replay information in #184. I'll add it back. Here's a minimal test case:
```py
def test_replay_map_data():
    pyro.set_rng_seed(0)
    data = list(range(10))

    def model():
        return pyro.map_data('mapdata', data, lambda i, x: i, batch_size=2)

    traced_model = poutine.trace(model)
    expected = traced_model()
    actual = poutine.replay(model, traced_model.trace)()
    assert actual == expected
```
@null-a Thanks for providing such an easily reproducible error!

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -10,6 +10,7 @@
 import pyro
 import pyro.poutine as poutine
 
+from pyro.distributions.subsample import Subsample
 from pyro.params import param_with_module_name
 from pyro.params.param_store import ParamStoreDict
 from pyro.poutine import LambdaPoutine, condition, do  # noqa: F401
@@ -149,7 +150,7 @@ def iarange(name, size, subsample_size=0):
         yield Variable(torch.LongTensor(list(range(size))))
         return
 
-    subsample = Variable(torch.randperm(size)[0:subsample_size])
+    subsample = sample(name, Subsample(size, subsample_size))
     if len(_PYRO_STACK) == 0:
         yield subsample
     else:
diff --git a/pyro/distributions/__init__.py b/pyro/distributions/__init__.py
--- a/pyro/distributions/__init__.py
+++ b/pyro/distributions/__init__.py
@@ -16,6 +16,7 @@
 from pyro.distributions.normal_chol import NormalChol
 from pyro.distributions.poisson import Poisson
 from pyro.distributions.uniform import Uniform
+from pyro.distributions.subsample import Subsample
 
 # function aliases
 diagnormal = DiagNormal()
@@ -33,3 +34,4 @@
 uniform = Uniform()
 dirichlet = Dirichlet()
 cauchy = Cauchy()
+subsample = Subsample()
diff --git a/pyro/distributions/subsample.py b/pyro/distributions/subsample.py
new file mode 100644
--- /dev/null
+++ b/pyro/distributions/subsample.py
@@ -0,0 +1,61 @@
+import torch
+from torch.autograd import Variable
+
+from pyro.distributions.distribution import Distribution
+from pyro.util import log_gamma
+
+
+class Subsample(Distribution):
+    """
+    Randomly select a subsample of a range of indices.
+
+    :param int size: the size of the range to subsample from
+    :param int subsample_size: the size of the returned subsample
+    :returns: a random subsample of `range(size)`
+    :rtype: torch.autograd.Variable of torch.LongTensor
+    """
+    reparameterized = False
+
+    def _sanitize_inputs(self, size, subsample_size):
+        if size is None:
+            size = self.size
+        if subsample_size is None:
+            subsample_size = self.subsample_size
+        return size, subsample_size
+
+    def __init__(self, size=None, subsample_size=None):
+        """
+        :param int size: the size of the range to subsample from
+        :param int subsample_size: the size of the returned subsample
+        """
+        self.size = size
+        self.subsample_size = subsample_size
+
+    def sample(self, size=None, subsample_size=None):
+        """
+        :param int size: the size of the range to subsample from
+        :param int subsample_size: the size of the returned subsample
+        :returns: a random subsample of `range(size)`
+        :rtype: torch.autograd.Variable of torch.LongTensor
+        """
+        size, subsample_size = self._sanitize_inputs(size, subsample_size)
+        if subsample_size > size:
+            raise ValueError("expected size < subsample_size, actual {} vs {}".format(
+                size, subsample_size))
+        return Variable(torch.randperm(size)[:subsample_size])
+
+    def batch_log_pdf(self, x, size=None, subsample_size=None):
+        """
+        :param int size: the size of the range to subsample from
+        :param int subsample_size: the size of the returned subsample
+        :param torch.autograd.Variable x: a subsample of `range(size)`
+        :return: log probability density
+        :rtype: torch.autograd.Variable
+        """
+        size, subsample_size = self._sanitize_inputs(size, subsample_size)
+        if subsample_size is None:
+            subsample_size = x.size(-1)
+        elif subsample_size != x.size(-1):
+            raise ValueError("subsample does not match subsample_size")
+        size = Variable(torch.Tensor([size]))
+        return log_gamma(1 + size) - log_gamma(1 + size - subsample_size)


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
@@ -1,3 +1,5 @@
+from __future__ import print_function
+
 import torch
 from torch.autograd import Variable
 import pytest
@@ -9,9 +11,8 @@
 from pyro.infer import SVI
 from tests.common import assert_equal
 
-pytestmark = pytest.mark.stage("integration", "integration_batch_1")
-
 
+@pytest.mark.stage("integration", "integration_batch_1")
 @pytest.mark.init(rng_seed=161)
 @pytest.mark.parametrize("batch_size", [3, 5, 7, 8, 0])
 @pytest.mark.parametrize("map_type", ["tensor", "list"])
@@ -155,5 +156,31 @@ def model(means, stds):
 
     tr = poutine.trace(model).get_trace(means, stds)
     for name in tr.nodes.keys():
-        if tr.nodes[name]["type"] == "sample":
+        if tr.nodes[name]["type"] == "sample" and name.startswith("x_"):
             assert tr.nodes[name]["scale"] == 4.0 * 2.0
+
+
+def test_replay_iarange():
+    pyro.set_rng_seed(0)
+
+    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 test_replay_map_data():
+    pyro.set_rng_seed(0)
+    data = list(range(10))
+
+    def model():
+        return pyro.map_data('mapdata', data, lambda i, x: i, batch_size=2)
+
+    traced_model = poutine.trace(model)
+    expected = traced_model()
+    actual = poutine.replay(model, traced_model.trace)()
+    assert actual == expected