Repository: pyro-ppl/pyro

Issue:
Create a poutine or messenger to configure site['infer']
It would be nice to have a poutine or messenger that could set `site['infer']` properties. Using this we could e.g. simplify the `enum_discrete` logic to first set `site['infer']['enumerate'] = 'sequential'` for all discrete sites, then use existing logic without `site_is_discrete`.


Discussion:
@eb8680 Would you be up to implement this? I'm not sure how.
@fritzo sure, I can do that

Resolving patch:
diff --git a/pyro/__init__.py b/pyro/__init__.py
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -327,6 +327,7 @@ def param(name, *args, **kwargs):
             "name": name,
             "args": args,
             "kwargs": kwargs,
+            "infer": {},
             "scale": 1.0,
             "cond_indep_stack": [],
             "value": None,
diff --git a/pyro/poutine/__init__.py b/pyro/poutine/__init__.py
--- a/pyro/poutine/__init__.py
+++ b/pyro/poutine/__init__.py
@@ -9,6 +9,7 @@
 from .condition_poutine import ConditionPoutine
 from .escape_poutine import EscapePoutine
 from .indep_poutine import IndepMessenger  # noqa: F401
+from .infer_config_poutine import InferConfigPoutine
 from .lift_poutine import LiftPoutine
 from .poutine import _PYRO_STACK, Poutine  # noqa: F401
 from .replay_poutine import ReplayPoutine
@@ -123,6 +124,20 @@ def condition(fn, data):
     return ConditionPoutine(fn, data=data)
 
 
+def infer_config(fn, config_fn):
+    """
+    :param fn: a stochastic function (callable containing pyro primitive calls)
+    :param config_fn: a callable taking a site and returning an infer dict
+
+    Alias for :class:`~pyro.poutine.infer_config_poutine.InferConfigPoutine` constructor.
+
+    Given a callable that contains Pyro primitive calls
+    and a callable taking a trace site and returning a dictionary,
+    updates the value of the infer kwarg at a sample site to config_fn(site)
+    """
+    return InferConfigPoutine(fn, config_fn)
+
+
 def indep(name, vectorized):
     """
     :param str name: a name for subsample sites
diff --git a/pyro/poutine/infer_config_poutine.py b/pyro/poutine/infer_config_poutine.py
new file mode 100644
--- /dev/null
+++ b/pyro/poutine/infer_config_poutine.py
@@ -0,0 +1,60 @@
+from __future__ import absolute_import, division, print_function
+
+from .poutine import Messenger, Poutine
+
+
+class InferConfigMessenger(Messenger):
+    """
+    Modifies contents of the infer kwarg at sample sites
+    """
+    def __init__(self, config_fn):
+        """
+        :param config_fn: a callable taking a site and returning an infer dict
+
+        Constructor. Doesn't do much, just stores the stochastic function
+        and the config_fn.
+        """
+        super(InferConfigMessenger, self).__init__()
+        self.config_fn = config_fn
+
+    def _pyro_sample(self, msg):
+        """
+        :param msg: current message at a trace site.
+
+        If self.config_fn is not None, calls self.config_fn on msg
+        and stores the result in msg["infer"].
+
+        Otherwise, implements default sampling behavior
+        with no additional effects.
+        """
+        msg["infer"].update(self.config_fn(msg))
+        return None
+
+    def _pyro_param(self, msg):
+        """
+        :param msg: current message at a trace site.
+
+        If self.config_fn is not None, calls self.config_fn on msg
+        and stores the result in msg["infer"].
+
+        Otherwise, implements default param behavior
+        with no additional effects.
+        """
+        msg["infer"].update(self.config_fn(msg))
+        return None
+
+
+class InferConfigPoutine(Poutine):
+    """
+    Modifies contents of the infer kwarg at sample sites
+    """
+    def __init__(self, fn, config_fn):
+        """
+        :param fn: a stochastic function (callable containing pyro primitive calls)
+        :param config_fn: a callable taking a site and returning an infer dict
+
+        Constructor. Doesn't do much, just stores the stochastic function
+        and the config_fn.
+        """
+        super(InferConfigPoutine, self).__init__(
+            InferConfigMessenger(config_fn=config_fn), fn)


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
@@ -592,3 +592,31 @@ def test_trace_compose(self):
                 assert False
             except NonlocalExit:
                 assert "x" not in tem.trace
+
+
+class InferConfigPoutineTests(TestCase):
+    def setUp(self):
+        def model():
+            pyro.param("p", Variable(torch.zeros(1), requires_grad=True))
+            pyro.sample("a", Bernoulli(Variable(torch.Tensor([0.5]))),
+                        infer={"enumerate": "parallel"})
+            pyro.sample("b", Bernoulli(Variable(torch.Tensor([0.5]))))
+
+        self.model = model
+
+        def config_fn(site):
+            if site["type"] == "sample":
+                return {"blah": True}
+            else:
+                return {}
+
+        self.config_fn = config_fn
+
+    def test_infer_config_sample(self):
+        cfg_model = poutine.infer_config(self.model, self.config_fn)
+
+        tr = poutine.trace(cfg_model).get_trace()
+
+        assert tr.nodes["a"]["infer"] == {"enumerate": "parallel", "blah": True}
+        assert tr.nodes["b"]["infer"] == {"blah": True}
+        assert tr.nodes["p"]["infer"] == {}