Repository: pyro-ppl/pyro

Issue:
Create mini-version of Pyro for didactic purposes
The Architecture of Open Source Applications project hosts tiny educational versions of many pieces of software https://github.com/aosabook/500lines (recommended reading!). I think it would be helpful to have a similar examples/mini_pyro.py to demonstrate how e.g. how Pyro's effects library works.

## Tasks

- [x] #1553 `trace`, `replay`, `SVI`, `elbo`, and a simple example
- [x] #1578 rendered version on pyro.ai
- [x] #1583 comments throughout the file
- [x] #1589 `plate` with broadcasting
- [x] #1599 `block` (@eb8680)

---- triaged ----
- ~~enumeration~~
- ~~`markov`~~

## Budget

All of pyro plus comments plus an example should fit into:
- 500 lines of code
- ~~1 single python file~~ 2 python files = 1 for example + 1 for library


Discussion:

Resolving patch:
diff --git a/pyro/contrib/minipyro.py b/pyro/contrib/minipyro.py
--- a/pyro/contrib/minipyro.py
+++ b/pyro/contrib/minipyro.py
@@ -18,7 +18,7 @@
 import torch
 
 # Pyro keeps track of two kinds of global state:
-# i)  The Poutine stack, which enables non-standard interpretations of
+# i)  The effect handler stack, which enables non-standard interpretations of
 #     Pyro primitives like sample();
 #     See http://docs.pyro.ai/en/0.3.0-release/poutine.html
 # ii) Trainable parameters in the Pyro ParamStore;
@@ -32,10 +32,13 @@ def get_param_store():
     return PARAM_STORE
 
 
+# The base effect handler class (called Messenger here for consistency with Pyro).
 class Messenger(object):
     def __init__(self, fn=None):
         self.fn = fn
 
+    # Effect handlers push themselves onto the PYRO_STACK.
+    # Handlers earlier in the PYRO_STACK are applied first.
     def __enter__(self):
         PYRO_STACK.append(self)
 
@@ -54,13 +57,19 @@ def __call__(self, *args, **kwargs):
             return self.fn(*args, **kwargs)
 
 
+# A first useful example of an effect handler.
+# trace records the inputs and outputs of any primitive site it encloses,
+# and returns a dictionary containing that data to the user.
 class trace(Messenger):
     def __enter__(self):
         super(trace, self).__enter__()
         self.trace = OrderedDict()
         return self.trace
 
+    # trace illustrates why we need postprocess_message in addition to process_message:
+    # We only want to record a value after all other effects have been applied
     def postprocess_message(self, msg):
+        assert msg["name"] not in self.trace, "all sites must have unique names"
         self.trace[msg["name"]] = msg.copy()
 
     def get_trace(self, *args, **kwargs):
@@ -68,6 +77,11 @@ def get_trace(self, *args, **kwargs):
         return self.trace
 
 
+# A second example of an effect handler for setting the value at a sample site.
+# This illustrates why effect handlers are a useful PPL implementation technique:
+# We can compose trace and replay to replace values but preserve distributions,
+# allowing us to compute the joint probability density of samples under a model.
+# See the definition of elbo(...) below for an example of this pattern.
 class replay(Messenger):
     def __init__(self, fn, guide_trace):
         self.guide_trace = guide_trace
@@ -78,6 +92,19 @@ def process_message(self, msg):
             msg["value"] = self.guide_trace[msg["name"]]["value"]
 
 
+# block allows the selective application of effect handlers to different parts of a model.
+# Sites hidden by block will only have the handlers below block on the PYRO_STACK applied,
+# allowing inference or other effectful computations to be nested inside models.
+class block(Messenger):
+    def __init__(self, fn=None, hide_fn=lambda msg: True):
+        self.hide_fn = hide_fn
+        super(block, self).__init__(fn)
+
+    def process_message(self, msg):
+        if self.hide_fn(msg):
+            msg["stop"] = True
+
+
 # This limited implementation of PlateMessenger only implements broadcasting.
 class PlateMessenger(Messenger):
     def __init__(self, fn, size, dim):
@@ -98,31 +125,36 @@ def __iter__(self):
         return range(self.size)
 
 
-def sample(name, fn, obs=None):
-    if not PYRO_STACK:
-        return fn()
-    msg = dict(type="sample", name=name, fn=fn, value=obs)
-    for handler in reversed(PYRO_STACK):
+def apply_stack(msg):
+    for pointer, handler in enumerate(reversed(PYRO_STACK)):
         handler.process_message(msg)
+        if msg.get("stop"):
+            break
     if msg["value"] is None:
-        msg["value"] = fn()
-    for handler in PYRO_STACK:
+        msg["value"] = msg["fn"](*msg["args"])
+    for handler in PYRO_STACK[-pointer-1:]:
         handler.postprocess_message(msg)
+    return msg
+
+
+def sample(name, fn, obs=None):
+    if not PYRO_STACK:
+        return fn()
+    msg = apply_stack(dict(type="sample", name=name, fn=fn, args=(), value=obs))
     return msg["value"]
 
 
+# param is an effectful version of PARAM_STORE.setdefault
 def param(name, init_value=None):
-    value = PARAM_STORE.setdefault(name, init_value)
-    value.requires_grad_()
-    if not PYRO_STACK:
+
+    def fn(init_value):
+        value = PARAM_STORE.setdefault(name, init_value)
+        value.requires_grad_()
         return value
-    msg = dict(type="param", name=name, value=None)
-    for handler in reversed(PYRO_STACK):
-        handler.process_message(msg)
-    if msg["value"] is None:
-        msg["value"] = value
-    for handler in PYRO_STACK:
-        handler.postprocess_message(msg)
+
+    if not PYRO_STACK:
+        return fn(init_value)
+    msg = apply_stack(dict(type="param", name=name, fn=fn, args=(init_value,), value=None))
     return msg["value"]
 
 
@@ -172,13 +204,13 @@ def step(self, *args, **kwargs):
         # we can record all the parameters that are encountered. Note that
         # further tracing occurs inside of `loss`.
         with trace() as param_capture:
-            loss = self.loss(self.model, self.guide, *args, **kwargs)
+            # We use block here to allow tracing to record parameters only.
+            with block(hide_fn=lambda msg: msg["type"] == "sample"):
+                loss = self.loss(self.model, self.guide, *args, **kwargs)
         # Differentiate the loss.
         loss.backward()
         # Grab all the parameters from the trace.
-        params = [site["value"]
-                  for site in param_capture.values()
-                  if site["type"] == "param"]
+        params = [site["value"] for site in param_capture.values()]
         # Take a step w.r.t. each parameter in params.
         self.optim(params)
         # Zero out the gradients so that they don't accumulate.
diff --git a/pyro/poutine/messenger.py b/pyro/poutine/messenger.py
--- a/pyro/poutine/messenger.py
+++ b/pyro/poutine/messenger.py
@@ -45,7 +45,7 @@ def __enter__(self):
         if not (self in _PYRO_STACK):
             # if this poutine is not already installed,
             # put it on the bottom of the stack.
-            _PYRO_STACK.insert(0, self)
+            _PYRO_STACK.append(self)
 
             # necessary to return self because the return value of __enter__
             # is bound to VAR in with EXPR as VAR.
@@ -88,8 +88,8 @@ def __exit__(self, exc_type, exc_value, traceback):
             # this poutine should be on the bottom of the stack.
             # If so, remove it from the stack.
             # if not, raise a ValueError because something really weird happened.
-            if _PYRO_STACK[0] == self:
-                _PYRO_STACK.pop(0)
+            if _PYRO_STACK[-1] == self:
+                _PYRO_STACK.pop()
             else:
                 # should never get here, but just in case...
                 raise ValueError("This Messenger is not on the bottom of the stack")
@@ -100,8 +100,8 @@ def __exit__(self, exc_type, exc_value, traceback):
             # then remove it and everything below it in the stack.
             if self in _PYRO_STACK:
                 loc = _PYRO_STACK.index(self)
-                for i in range(0, loc + 1):
-                    _PYRO_STACK.pop(0)
+                for i in range(loc, len(_PYRO_STACK)):
+                    _PYRO_STACK.pop()
 
     def _reset(self):
         pass
diff --git a/pyro/poutine/runtime.py b/pyro/poutine/runtime.py
--- a/pyro/poutine/runtime.py
+++ b/pyro/poutine/runtime.py
@@ -137,7 +137,7 @@ def reset_stack(self):
         Reset the state of the frames remaining in the stack.
         Necessary for multiple re-executions in poutine.queue.
         """
-        for frame in _PYRO_STACK:
+        for frame in reversed(_PYRO_STACK):
             frame._reset()
             if type(frame).__name__ == "BlockMessenger" and frame.hide_fn(self.site):
                 break
@@ -181,11 +181,11 @@ def apply_stack(initial_msg):
     # msg is used to pass information up and down the stack
     msg = initial_msg
 
-    counter = 0
+    pointer = 0
     # go until time to stop?
-    for frame in stack:
+    for frame in reversed(stack):
 
-        counter = counter + 1
+        pointer = pointer + 1
 
         frame._process_message(msg)
 
@@ -194,7 +194,7 @@ def apply_stack(initial_msg):
 
     default_process_message(msg)
 
-    for frame in reversed(stack[0:counter]):
+    for frame in stack[-pointer:]:  # reversed(stack[0:pointer])
         frame._postprocess_message(msg)
 
     cont = msg["continuation"]


Test patch: