Repository: pyro-ppl/pyro

Issue:
Move distributions.Subsample into pyro.util
Also
- set `log_pdf` to zero
- do not support batching
- remove references in `docs/source`

## Why?

`Subsample` is an internal distribution with special handling: it should never be batched (since it is used to define batching), and its `log_pdf` should not be used to change the `log_pdf` (i.e. `Subsample.log_pdf` must be zero so that `map_data` can provide an unbiased estimator). 

## How?

 The final implementation should be ~10 lines plus a single docstring:
```py
class Subsample(distribution):
    """
    Internal use only. Randomly subsample a linear range of indices.
    """
    def __init__(self, size, batch_size):
        assert 0 <= batch_size <= size
        self._size = size
        self._batch_size = batch_size

    def sample(self):
        return Variable(...)

    def log_pdf(self, *args, **kwargs):
        return ng_zeros(1)
```


Discussion:

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -10,7 +10,7 @@
 import pyro
 import pyro.poutine as poutine
 
-from pyro.distributions.subsample import Subsample
+from pyro.distributions.distribution import Distribution
 from pyro.params import param_with_module_name
 from pyro.params.param_store import ParamStoreDict
 from pyro.poutine import LambdaPoutine, condition, do  # noqa: F401
@@ -116,6 +116,35 @@ def observe(name, fn, obs, *args, **kwargs):
     return sample(name, fn, *args, **kwargs)
 
 
+class _Subsample(Distribution):
+    """
+    Randomly select a subsample of a range of indices.
+
+    Internal use only. This should only be used by `iarange`.
+    """
+
+    def __init__(self, size, subsample_size):
+        """
+        :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):
+        """
+        :returns: a random subsample of `range(size)`
+        :rtype: torch.autograd.Variable of torch.LongTensor
+        """
+        assert 0 <= self.subsample_size <= self.size
+        return Variable(torch.randperm(self.size)[:self.subsample_size])
+
+    def batch_log_pdf(self, x):
+        # This is zero so that iarange can provide an unbiased estimate of
+        # the non-subsampled batch_log_pdf.
+        return Variable(torch.zeros(0))
+
+
 @contextlib.contextmanager
 def iarange(name, size, subsample_size=0):
     """
@@ -151,7 +180,7 @@ def iarange(name, size, subsample_size=0):
         yield Variable(torch.LongTensor(list(range(size))))
         return
 
-    subsample = sample(name, Subsample(size, 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,7 +16,6 @@
 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()
@@ -34,4 +33,3 @@
 uniform = Uniform()
 dirichlet = Dirichlet()
 cauchy = Cauchy()
-subsample = Subsample()
diff --git a/pyro/distributions/subsample.py b/pyro/distributions/subsample.py
deleted file mode 100644
--- a/pyro/distributions/subsample.py
+++ /dev/null
@@ -1,60 +0,0 @@
-import torch
-from torch.autograd import Variable
-
-from pyro.distributions.distribution import Distribution
-from pyro.distributions.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
-    """
-
-    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: