Repository: pyro-ppl/pyro

Issue:
poutine.block is thwarted by _DIM_ALLOCATOR global state
## Desired behavior

Any submodel can be run inside a Pyro model by using `poutine.block`. These submodels should have no side effects.

## Actual behavior

Submodels may modify the `_DIM_ALLOCATOR`. This is problematic in composite ADVI where two ADVI strategies gather different traces which allocate duplicate dimensions for a shared iarange. See #1002 .


Discussion:
@eb8680 I've gotten myself into a mess of side effects here. Can you help me get out?
Hmm it turns out our issue was unrelated, but we saw tests failing due to `_DIM_ALLOCATOR` global state that had been left dirty by previous failing tests. Not sure whether this issue is really a bug...
> we saw tests failing due to _DIM_ALLOCATOR global state that had been left dirty by previous failing tests.

Strange, this should be handled by `iarange.__exit__`.  I thought I knew what was causing this but I think I was wrong, I'll investigate further.
To give more context: `_DIM_ALLOCATOR` was left in a dirty state after an error occurred *inside* the `iarange.__enter__()` method. Here is the stack trace that led to the dirty state:

<details>

```
______________________________________________ test_discrete_parallel[ADVIMultivariateNormal] ______________________________________________

continuous_class = <class 'pyro.infer.advi.ADVIMultivariateNormal'>

    @pytest.mark.parametrize("continuous_class", [ADVIMultivariateNormal, ADVIDiagonalNormal])
    def test_discrete_parallel(continuous_class):
        K = 2
        data = torch.tensor([0., 1., 10., 11., 12.])

        def model(data):
            weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
            locs = pyro.sample('locs', dist.Normal(0, 10).reshape([K]))
            scale = pyro.sample('scale', dist.LogNormal(0, 1))

            with pyro.iarange('data'):
                weights = weights.expand(torch.Size((len(data),)) + weights.shape)
                assignment = pyro.sample('assignment', dist.Categorical(weights))
                pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)

        advi = ADVIMaster(model)
        advi.add(ADVIDiscreteParallel(poutine.block(model, expose=["assignment"])))
        advi.add(continuous_class(poutine.block(model, hide=["assignment"])))
        advi.setup_prototype(data)

        elbo = ELBO.make(enum_discrete=True)
>       loss = elbo.loss_and_grads(advi.model, advi.guide, data)

tests/infer/test_advi.py:125:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
pyro/infer/traceenum_elbo.py:113: in loss_and_grads
    for model_trace, guide_trace in self._get_traces(model, guide, *args, **kwargs):
pyro/infer/traceenum_elbo.py:61: in _get_traces
    for guide_trace in iter_discrete_traces("flat", guide, *args, **kwargs):
pyro/infer/enum.py:48: in iter_discrete_traces
    yield traced_fn.get_trace(*args, **kwargs)
pyro/poutine/trace_poutine.py:182: in get_trace
    self(*args, **kwargs)
pyro/poutine/trace_poutine.py:170: in __call__
    ret = super(TracePoutine, self).__call__(*args, **kwargs)
pyro/poutine/poutine.py:147: in __call__
    return self.fn(*args, **kwargs)
pyro/poutine/__init__.py:218: in _fn
    return ftr(*args, **kwargs)
pyro/poutine/trace_poutine.py:170: in __call__
    ret = super(TracePoutine, self).__call__(*args, **kwargs)
pyro/poutine/poutine.py:147: in __call__
    return self.fn(*args, **kwargs)
pyro/poutine/poutine.py:147: in __call__
    return self.fn(*args, **kwargs)
pyro/poutine/poutine.py:147: in __call__
    return self.fn(*args, **kwargs)
pyro/poutine/poutine.py:147: in __call__
    return self.fn(*args, **kwargs)
pyro/infer/advi.py:74: in guide
    part.guide(*args, **kwargs)
pyro/infer/advi.py:417: in guide
    stack.enter_context(iaranges[frame.name])
../../../miniconda2/envs/pytorch-dev/lib/python2.7/site-packages/contextlib2.py:380: in enter_context
    result = _cm_type.__enter__(cm)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <pyro.iarange object at 0x107be9210>

    def __enter__(self):
        self._wrapped = am_i_wrapped()
        self.dim = _DIM_ALLOCATOR.allocate(self.name, self.dim)
        if self._wrapped:
>           self._scale_poutine = poutine.ScaleMessenger(self.size / self.subsample_size)
E           ZeroDivisionError: division by zero

pyro/__init__.py:251: ZeroDivisionError
```
</details>
Seems like the solution is to wrap the rest of `iarange.__enter__` in a `try`/`except` and call `_DIM_ALLOCATOR.free` before propagating the exception further.

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -248,10 +248,14 @@ def __enter__(self):
         self._wrapped = am_i_wrapped()
         self.dim = _DIM_ALLOCATOR.allocate(self.name, self.dim)
         if self._wrapped:
-            self._scale_poutine = poutine.ScaleMessenger(self.size / self.subsample_size)
-            self._indep_poutine = poutine.IndepMessenger(self.name, size=self.subsample_size, dim=self.dim)
-            self._scale_poutine.__enter__()
-            self._indep_poutine.__enter__()
+            try:
+                self._scale_poutine = poutine.ScaleMessenger(self.size / self.subsample_size)
+                self._indep_poutine = poutine.IndepMessenger(self.name, size=self.subsample_size, dim=self.dim)
+                self._scale_poutine.__enter__()
+                self._indep_poutine.__enter__()
+            except BaseException:
+                _DIM_ALLOCATOR.free(self.name, self.dim)
+                raise
         return self.subsample
 
     def __exit__(self, exc_type, exc_value, traceback):


Test patch:
diff --git a/tests/poutine/test_poutines.py b/tests/poutine/test_poutines.py
--- a/tests/poutine/test_poutines.py
+++ b/tests/poutine/test_poutines.py
@@ -6,13 +6,14 @@
 import pytest
 import torch
 import torch.nn as nn
+from six.moves.queue import Queue
 
 import pyro
 import pyro.distributions as dist
 import pyro.poutine as poutine
 from pyro.distributions import Bernoulli, Categorical, Normal
-from pyro.poutine.util import all_escape, discrete_escape, NonlocalExit
-from six.moves.queue import Queue
+from pyro.poutine.indep_poutine import _DIM_ALLOCATOR
+from pyro.poutine.util import NonlocalExit, all_escape, discrete_escape
 from tests.common import assert_equal
 
 
@@ -669,3 +670,14 @@ def model():
         actual_shape = log_prob.shape
         expected_shape = (2,) * depth + (3,) + (2,) * depth + (1,) * first_available_dim
         assert actual_shape == expected_shape, 'error on iteration {}'.format(i)
+
+
+def test_iarange_error_on_enter():
+    def model():
+        with pyro.iarange('foo', 0):
+            pass
+
+    assert len(_DIM_ALLOCATOR._stack) == 0
+    with pytest.raises(ZeroDivisionError):
+        poutine.trace(model)()
+    assert len(_DIM_ALLOCATOR._stack) == 0, 'stack was not cleaned on error'