Repository: pyro-ppl/pyro

Issue:
minipyro documentation clarification.
I really like that minipyro is in here for those interested in the internals of pyro. Super awesome!

Some of the wording in the module, particularly regarding `PYRO_STACK`, is confusing to me, though.

For example:
https://github.com/pyro-ppl/pyro/blob/dc7b96471b3e2dd32f6f39eb2e10ccedc2a4a8a3/pyro/contrib/minipyro.py#L48

If we run:
```python
trace(replay(model, guide_trace)).get_trace(...)
```

We end up (somewhere during the execution of the above line) with a `PYRO_STACK` with the first item referencing `trace`, and the second item referencing `replay`. But when we hit `apply_stack`, because of the reverse order in which the handlers are applied, `replay` is applied first. So, the comments should read, `# Handlers *later* in the PYRO_STACK are applied first.`, correct?

Similarly, in
https://github.com/pyro-ppl/pyro/blob/dc7b96471b3e2dd32f6f39eb2e10ccedc2a4a8a3/pyro/contrib/minipyro.py#L164

Shouldn't 
> When a Messenger sets the "stop" field of a message,
> it prevents any Messengers above it on the stack from being applied.

really be:
> When a Messenger sets the "stop" field of a message,
> it prevents any Messengers *below* it on the stack from being applied.

And
> A Messenger that sets msg["stop"] == True also prevents application
> of postprocess_message by Messengers above it on the stack
> via the pointer variable from the process_message loop

should really be:
> A Messenger that sets msg["stop"] == True also prevents application
> of postprocess_message by Messengers *below* it on the stack
> via the pointer variable from the process_message loop

I could be convinced to interpret all this otherwise, but I guess this was my default interpretation. Am I way off-base?


Discussion:
Hi @luiarthur, yes your reading seems to be the correct interpretation. We'd be happy to review a PR fixing or clarifying the comments!
Cool! I'll start a PR.

Resolving patch:
diff --git a/pyro/contrib/minipyro.py b/pyro/contrib/minipyro.py
--- a/pyro/contrib/minipyro.py
+++ b/pyro/contrib/minipyro.py
@@ -45,7 +45,7 @@ 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.
+    # Handlers later in the PYRO_STACK are applied first.
     def __enter__(self):
         PYRO_STACK.append(self)
 
@@ -162,17 +162,19 @@ def __iter__(self):
 # apply_stack is called by pyro.sample and pyro.param.
 # It is responsible for applying each Messenger to each effectful operation.
 def apply_stack(msg):
+    # PYRO_STACK is reversed so that effect handlers higher
+    # in the stack are first applied.
     for pointer, handler in enumerate(reversed(PYRO_STACK)):
         handler.process_message(msg)
         # When a Messenger sets the "stop" field of a message,
-        # it prevents any Messengers above it on the stack from being applied.
+        # it prevents any Messengers below it on the stack from being applied.
         if msg.get("stop"):
             break
     if msg["value"] is None:
         msg["value"] = msg["fn"](*msg["args"])
 
     # A Messenger that sets msg["stop"] == True also prevents application
-    # of postprocess_message by Messengers above it on the stack
+    # of postprocess_message by Messengers below it on the stack
     # via the pointer variable from the process_message loop
     for handler in PYRO_STACK[-pointer - 1 :]:
         handler.postprocess_message(msg)


Test patch: