Exclude seed handler from the PYRO_STACK

pyro/contrib/minipyro.py

@@ -116,14 +116,23 @@ def __init__(self, fn=None, rng_seed=None):
         super(seed, self).__init__(fn)
 
     def __enter__(self):
+        self.old_state = {'torch': torch.get_rng_state(), 'random': random.getstate()}
         torch.manual_seed(self.rng_seed)
         random.seed(self.rng_seed)
         try:
             import numpy as np
             np.random.seed(self.rng_seed)
+            self.old_state['numpy'] = np.random.get_state()
         except ImportError:
             pass
 
+    def __exit__(self, type, value, traceback):
+        torch.set_rng_state(self.old_state['torch'])
+        random.setstate(self.old_state['random'])
+        if 'numpy' in self.old_state:
+            import numpy as np
+            np.random.set_state(self.old_state['numpy'])
+
 
 # This limited implementation of PlateMessenger only implements broadcasting.
 class PlateMessenger(Messenger):
@@ -229,7 +238,9 @@ def fn(init_value, constraint):
 
 
 # boilerplate to match the syntax of actual pyro.plate:
-def plate(name, size, dim):
+def plate(name, size, dim=None):
+    if dim is None:
+        raise NotImplementedError("minipyro.plate requires a dim arg")
     return PlateMessenger(fn=None, size=size, dim=dim)
 
 

pyro/poutine/handlers.py

@@ -50,7 +50,8 @@
 
 from pyro.poutine import util
 from pyro.poutine.messenger import Messenger
-from pyro.util import set_rng_seed
+from pyro.util import get_rng_state, set_rng_seed, set_rng_state
+
 from .block_messenger import BlockMessenger
 from .broadcast_messenger import BroadcastMessenger
 from .condition_messenger import ConditionMessenger
@@ -67,7 +68,6 @@
 from .trace_messenger import TraceMessenger
 from .uncondition_messenger import UnconditionMessenger
 
-
 ############################################
 # Begin primitive operations
 ############################################
@@ -508,8 +508,11 @@ def __init__(self, rng_seed):
         super(_SeedMessenger, self).__init__()
 
     def __enter__(self):
+        self.old_state = get_rng_state()
         set_rng_seed(self.rng_seed)
-        super(_SeedMessenger, self).__enter__()
+
+    def __exit__(self, type, value, traceback):
+        set_rng_state(self.old_state)
 
 
 def seed(fn=None, rng_seed=None):

pyro/util.py

@@ -28,6 +28,24 @@ def set_rng_seed(rng_seed):
         pass
 
 
+def get_rng_state():
+    state = {'torch': torch.get_rng_state(), 'random': random.getstate()}
+    try:
+        import numpy as np
+        state['numpy'] = np.random.get_state()
+    except ImportError:
+        pass
+    return state
+
+
+def set_rng_state(state):
+    torch.set_rng_state(state['torch'])
+    random.setstate(state['random'])
+    if 'numpy' in state:
+        import numpy as np
+        np.random.set_state(state['numpy'])
+
+
 def torch_isnan(x):
     """
     A convenient function to check if a Tensor contains any nan; also works with numbers

tests/test_generic.py

@@ -2,6 +2,7 @@
 
 from pyro.generic import handlers, infer, pyro, pyro_backend
 from pyro.generic.testing import MODELS
+from tests.common import assert_close, xfail_if_not_implemented
 
 pytestmark = pytest.mark.stage('unit')
 
@@ -30,19 +31,47 @@ def test_not_implemented(backend):
 
 @pytest.mark.parametrize('model', MODELS)
 @pytest.mark.parametrize('backend', ['minipyro', 'pyro'])
-@pytest.mark.xfail(reason='Not supported by backend.')
 def test_model_sample(model, backend):
-    with pyro_backend(backend), handlers.seed(rng_seed=2):
+    with pyro_backend(backend), handlers.seed(rng_seed=2), xfail_if_not_implemented():
         f = MODELS[model]()
         model, model_args = f['model'], f.get('model_args', ())
         model(*model_args)
 
 
 @pytest.mark.parametrize('model', MODELS)
 @pytest.mark.parametrize('backend', ['minipyro', 'pyro'])
-@pytest.mark.xfail(reason='Not supported by backend.')
+def test_rng_seed(model, backend):
+    with pyro_backend(backend), handlers.seed(rng_seed=2), xfail_if_not_implemented():
+        f = MODELS[model]()
+        model, model_args = f['model'], f.get('model_args', ())
+        with handlers.seed(rng_seed=0):
+            expected = model(*model_args)
+        with handlers.seed(rng_seed=0):
+            actual = model(*model_args)
+        assert_close(actual, expected)
+
+
+@pytest.mark.parametrize('model', MODELS)
+@pytest.mark.parametrize('backend', ['minipyro', 'pyro'])
+def test_rng_state(model, backend):
+    with pyro_backend(backend), handlers.seed(rng_seed=2), xfail_if_not_implemented():
+        f = MODELS[model]()
+        model, model_args = f['model'], f.get('model_args', ())
+        with handlers.seed(rng_seed=0):
+            model(*model_args)
+            expected = model(*model_args)
+        with handlers.seed(rng_seed=0):
+            model(*model_args)
+            with handlers.seed(rng_seed=0):
+                model(*model_args)
+            actual = model(*model_args)
+        assert_close(actual, expected)
+
+
+@pytest.mark.parametrize('model', MODELS)
+@pytest.mark.parametrize('backend', ['minipyro', 'pyro'])
 def test_trace_handler(model, backend):
-    with pyro_backend(backend), handlers.seed(rng_seed=2):
+    with pyro_backend(backend), handlers.seed(rng_seed=2), xfail_if_not_implemented():
         f = MODELS[model]()
         model, model_args, model_kwargs = f['model'], f.get('model_args', ()), f.get('model_kwargs', {})
         # should be implemented