Repository: pyro-ppl/numpyro

Issue:
Scan fails with empty PYRO_STACK
When `scan` is taken outside the context of a non-empty `PYRO_STACK`, `scan` produces an error.

The following code produces the error. I don't think it matters much what the scan actually does, but in this case it's interative multiplication of a vector by a matrix with the addition of an second vector at each step.

```python
def matrix_rollout(K, c_in):
    
    def iteration(c_prev, c_in):
        c_next = jnp.dot(c_prev, K) + c_in
        return c_next, (c_next,)
    
    _, (ys,) = scan(iteration, init=jnp.asarray([1.0, 0.0]), xs=c_in)

    return ys

matrix_rollout(
    jnp.asarray([[0.7, 0.3],
                 [0.3, 0.7]]),
    c_in=jnp.asarray([
        [1.0, 0.0],
    ])
```

Produces

```
---------------------------------------------------------------------------
UnboundLocalError                         Traceback (most recent call last)
<ipython-input-409-2df128ef5191> in <module>
     12     return ys
     13 
---> 14 matrix_rollout(
     15     jnp.asarray([[0.7, 0.3],
     16                  [0.3, 0.7]]),

<ipython-input-409-2df128ef5191> in matrix_rollout(K, c_in)
      8         return c_next, (c_next,)
      9 
---> 10     _, (ys,) = scan(iteration, init=jnp.asarray([1.0, 0.0]), xs=c_in)
     11 
     12     return ys

numpyro/contrib/control_flow/scan.py in scan(f, init, xs, length, reverse, history)
    397         (length, rng_key, carry), (pytree_trace, ys) = msg['value']
    398 
--> 399     if not msg["kwargs"].get("enum", False):
    400         for msg in pytree_trace.trace.values():
    401             apply_stack(msg)

UnboundLocalError: local variable 'msg' referenced before assignment
```

The relevant source code is [here](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/contrib/control_flow/scan.py#L468). I don't really understand how messengers work, but it looks like it might be an easy fix. I'd be happy to do it myself, but I'm not sure what all this messengers code does...


Discussion:
@justinrporter Thanks, I think you can return `ys, carry` at [this line](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/contrib/control_flow/scan.py#L453). Could you help me submit the fix?

Resolving patch:
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
--- a/numpyro/contrib/control_flow/scan.py
+++ b/numpyro/contrib/control_flow/scan.py
@@ -451,6 +451,7 @@ def g(*args, **kwargs):
         (length, rng_key, carry), (pytree_trace, ys) = scan_wrapper(
             f, init, xs, length=length, reverse=reverse
         )
+        return carry, ys
     else:
         # Otherwise, we initialize a message...
         initial_msg = {


Test patch:
diff --git a/test/contrib/test_control_flow.py b/test/contrib/test_control_flow.py
--- a/test/contrib/test_control_flow.py
+++ b/test/contrib/test_control_flow.py
@@ -112,3 +112,23 @@ def transition(state, i):
     actual_log_joint = potential_energy(fun_model, (T,), {}, fun_params)
     expected_log_joint = potential_energy(model, (T,), {}, params)
     assert_allclose(actual_log_joint, expected_log_joint)
+
+
+def test_scan_without_stack():
+    def multiply_and_add_repeatedly(K, c_in):
+        def iteration(c_prev, c_in):
+            c_next = jnp.dot(c_prev, K) + c_in
+            return c_next, (c_next,)
+
+        _, (ys,) = scan(iteration, init=jnp.asarray([1.0, 0.0]), xs=c_in)
+
+        return ys
+
+    result = multiply_and_add_repeatedly(
+        K=jnp.asarray([[0.7, 0.3], [0.3, 0.7]]), c_in=jnp.asarray([[1.0, 0.0]])
+    )
+
+    assert_allclose(
+        result,
+        [[1.7, 0.3]],
+    )