Repository: pyro-ppl/pyro

Issue:
Documentation for `iarange`
http://docs.pyro.ai/primitives.html#pyro.__init__.iarange

I had trouble understanding this documentation. Seems to be written for a maintainer as it refers to several things that are not in the docs: context manager, traces, etc. 





Discussion:
Thanks for pointing this out, this is definitely too obtuse for a user-facing construct.  We should fix the docstring and possibly write a dedicated tutorial about subsampling as requested by #67 .
it may be worth pointing out that all of this is discussed in the second svi tutorial. (at least from a user's perspective). which isn't to say the discussion couldn't be improved.
Yeah, tutorial is helpful, but ideally the docs shouldn't require it.

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -32,8 +32,8 @@ def clear_param_store():
 
 def sample(name, fn, *args, **kwargs):
     """
-    Calls the stochastic function `fn` with additional side-effects depending on `name` and the
-    enclosing context (e.g. an inference algorithm).
+    Calls the stochastic function `fn` with additional side-effects depending
+    on `name` and the enclosing context (e.g. an inference algorithm).
     See `Intro I <http://pyro.ai/examples/intro_part_i.html>`_ and
     `Intro II <http://pyro.ai/examples/intro_part_ii.html>`_ for a discussion.
 
@@ -160,16 +160,38 @@ def _subsample(name, size=None, subsample_size=None, subsample=None, use_cuda=No
 @contextlib.contextmanager
 def iarange(name, size=None, subsample_size=None, subsample=None, use_cuda=None):
     """
-    Context manager for ranges indexing iid variables, optionally subsampling.
-
-    WARNING: This is only correct if all computation is iid within the context.
-
-    By default `subsample_size=False` and this simply yields a
-    `torch.arange(0, size)`. If `0 < subsample_size <= size` this yields a
-    single random batch of indices of size `subsample_size` and scales all log
-    likelihood terms by `size/batch_size`, within this context.
-
-    :param str name: A name that will be used for this site in a Trace.
+    Context manager for conditionally independent ranges of variables.
+
+    ``iarange`` is similar to ``torch.arange`` in that it yields an array
+    of indices by which other tensors can be indexed. ``iarange`` differs from
+    ``torch.arange`` in that it also informs inference algorithms that the
+    variables being indexed are conditionally independent. To do this,
+    ``iarange`` is a provided as context manager rather than a function, and
+    users must guarantee that all computation within an ``iarange`` context
+    is conditionally independent::
+
+        with iarange("name", size) as ind:
+            # ...do conditionally independent stuff with ind...
+
+    Additionally, ``iarange`` can take advantage of the conditional
+    independence assumptions by subsampling the indices and informing inference
+    algorithms to scale various computed values. This is typically used to
+    subsample minibatches of data::
+
+        with iarange("data", len(data), subsample_size=100) as ind:
+            batch = data[ind]
+            assert len(batch) == 100
+
+    By default ``subsample_size=False`` and this simply yields a
+    ``torch.arange(0, size)``. If ``0 < subsample_size <= size`` this yields a
+    single random batch of indices of size ``subsample_size`` and scales all
+    log likelihood terms by ``size/batch_size``, within this context.
+
+    .. warning::  This is only correct if all computation is conditionally
+        independent within the context.
+
+    :param str name: A unique name to help inference algorithms match
+        ``iarange`` sites between models and guides.
     :param int size: Optional size of the collection being subsampled
         (like `stop` in builtin `range`).
     :param int subsample_size: Size of minibatches used in subsampling.
@@ -198,7 +220,8 @@ def iarange(name, size=None, subsample_size=None, subsample=None, use_cuda=None)
         >>> with iarange('data', 100, subsample=ind):
                 observe('obs', normal, data.index_select(0, ind), mu, sigma)
 
-    See `SVI Part II <http://pyro.ai/examples/svi_part_ii.html>`_ for an extended discussion.
+    See `SVI Part II <http://pyro.ai/examples/svi_part_ii.html>`_ for an
+    extended discussion.
     """
     subsample, scale = _subsample(name, size, subsample_size, subsample, use_cuda)
     if len(_PYRO_STACK) == 0:
@@ -211,17 +234,21 @@ def iarange(name, size=None, subsample_size=None, subsample=None, use_cuda=None)
 
 def irange(name, size, subsample_size=None, subsample=None, use_cuda=None):
     """
-    Non-vectorized version of `iarange`. See `iarange` for details.
+    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`.
-    :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()`.
+    :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``.
+    :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()``.
     :param bool use_cuda: Optional bool specifying whether to use cuda tensors
-        for `log_pdf`. Defaults to `torch.Tensor.is_cuda`.
-    :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
+        for internal ``log_pdf`` computations. Defaults to
+        ``torch.Tensor.is_cuda``.
+    :return: A generator yielding a sequence of integers.
 
     Examples::
 


Test patch: