Repository: pyro-ppl/pyro

Issue:
composable guides/autoguides
Hey team,

I was looking at the autoguides a bit and have the following request.

It would be cool to be able to define guides over some parts of the model (imagine I have custom inference strategies for particular components int he model) and let autoguides take care of the rest.

This would correspond to composing guides to explain a model and would require some mechanism that would test which variables are taken care of explicitly in the custom guide and which other ones need to be automagically guided.

Currently it feels as if there is some component missing to do this, as we are not attaching scopes to autoguides, but I may also have missed some discussions of course.

One might also consider having two different autoguides scoped over different parts of the model (i.e. a neural autoguide for local variables and an advi autoguide over globals), although I have not thought this particular case through and it may have issues with building proper inverse models and conditioning on previous samples and so on.

Any thoughts on this?

With Regards,

Theo


Discussion:
Interesting idea!  We should be able to update `AutoGuideList` to do a very simple version of this:
```python
guide = AutoGuideList(model)
guide.add(AutoDiagonalNormal(poutine.block(model, expose=[...])))
guide.add(custom_guide)  # assume custom_guide does not overlap with the AutoDiagonalNormal
...
svi = SVI(model, guide, optim, loss)
for batch in data:
    svi.step(batch)
```

You can also alternate updates for each guide with `poutine.condition`:
```python
guide1 = AutoDiagonalNormal(poutine.block(model, expose=[...]))
guide2 = custom_guide
optim1 = pyro.optim.Adam(...)
optim2 = pyro.optim.Adam(...)
while True:
    guide1_tr = poutine.trace(guide1).get_trace(...)
    SVI(poutine.condition(model, data=guide1_tr), guide2, optim2, loss).step(...)
    guide2_tr = poutine.trace(guide2).get_trace(...)
    SVI(poutine.condition(model, data=guide2_tr), guide1, optim1, loss).step(...)
````
Thanks Eli, this looks really nice and idiomatic with poutine magic.
I am really a fan of solution 1, as it seems like a very natural way to define composable guides in the sense I had in mind. Would be great to have that functionality.

Solution 2 also seems ok, but I prefer the elegance of 1, I do not necessarily want to force the alternating updates, but it is good to see how to do it as this may be useful in different contexts.


You can also write solution 1 directly, assuming there are no name collisions:
```python
guide_list = [AutoDiagonalNormal(poutine.block(model, ...)), custom_guide]
composite_guide = lambda *args: tuple(guide(*args) for guide in guide_list)

svi = SVI(model, composite_guide, optim, loss)

svi.step(...)
```

Resolving patch:
diff --git a/pyro/contrib/autoguide/__init__.py b/pyro/contrib/autoguide/__init__.py
--- a/pyro/contrib/autoguide/__init__.py
+++ b/pyro/contrib/autoguide/__init__.py
@@ -34,6 +34,7 @@ def model():
     from contextlib2 import ExitStack  # python 2
 
 __all__ = [
+    'AutoCallable',
     'AutoContinuous',
     'AutoDelta',
     'AutoDiagonalNormal',
@@ -157,9 +158,11 @@ def add(self, part):
         have been created by blocking the model to restrict to a subset of
         sample sites. No two parts should operate on any one sample site.
 
-        :param AutoGuide part: a partial guide to add
+        :param part: a partial guide to add
+        :type part: AutoGuide or callable
         """
-        assert isinstance(part, AutoGuide), type(part)
+        if not isinstance(part, AutoGuide):
+            part = AutoCallable(self.model, part)
         self.parts.append(part)
         assert part.master is None
         part.master = weakref.ref(self)
@@ -198,6 +201,45 @@ def median(self, *args, **kwargs):
         return result
 
 
+class AutoCallable(AutoGuide):
+    """
+    :class:`AutoGuide` wrapper for simple callable guides.
+
+    This is used internally for composing autoguides with custom user-defined
+    guides that are simple callables, e.g.::
+
+        def my_local_guide(*args, **kwargs):
+            ...
+
+        guide = AutoGuideList(model)
+        guide.add(AutoDelta(poutine.block(model, expose=['my_global_param']))
+        guide.add(my_local_guide)  # automatically wrapped in an AutoCallable
+
+    To specify a median callable, you can instead::
+
+        def my_local_median(*args, **kwargs)
+            ...
+
+        guide.add(AutoCallable(model, my_local_guide, my_local_median))
+
+    For more complex guides that need e.g. access to iaranges, users should
+    instead subclass ``AutoGuide``.
+
+    :param callable model: a Pyro model
+    :param callable guide: a Pyro guide (typically over only part of the model)
+    :param callable median: an optional callable returning a dict mapping
+        sample site name to computed median tensor.
+    """
+    def __init__(self, model, guide, median=lambda *args, **kwargs: {}):
+        super(AutoCallable, self).__init__(model, prefix="")
+        self._guide = guide
+        self.median = median
+
+    def __call__(self, *args, **kwargs):
+        result = self._guide(*args, **kwargs)
+        return {} if result is None else result
+
+
 class AutoDelta(AutoGuide):
     """
     This implementation of :class:`AutoGuide` uses Delta distributions to


Test patch:
diff --git a/tests/contrib/autoguide/test_advi.py b/tests/contrib/autoguide/test_advi.py
--- a/tests/contrib/autoguide/test_advi.py
+++ b/tests/contrib/autoguide/test_advi.py
@@ -3,11 +3,12 @@
 import numpy as np
 import pytest
 import torch
+from torch.distributions import constraints
 
 import pyro
 import pyro.distributions as dist
 import pyro.poutine as poutine
-from pyro.contrib.autoguide import (AutoDelta, AutoDiagonalNormal, AutoDiscreteParallel, AutoGuideList,
+from pyro.contrib.autoguide import (AutoCallable, AutoDelta, AutoDiagonalNormal, AutoDiscreteParallel, AutoGuideList,
                                     AutoLowRankMultivariateNormal, AutoMultivariateNormal)
 from pyro.infer import SVI, Trace_ELBO, TraceEnum_ELBO, TraceGraph_ELBO
 from pyro.optim import Adam
@@ -92,12 +93,28 @@ def auto_guide_list_x(model):
     return guide
 
 
+def auto_guide_callable(model):
+    def guide_x():
+        x_loc = pyro.param("x_loc", torch.tensor(1.))
+        x_scale = pyro.param("x_scale", torch.tensor(2.), constraint=constraints.positive)
+        pyro.sample("x", dist.Normal(x_loc, x_scale))
+
+    def median_x():
+        return {"x": pyro.param("x_loc", torch.tensor(1.))}
+
+    guide = AutoGuideList(model)
+    guide.add(AutoCallable(model, guide_x, median_x))
+    guide.add(AutoDiagonalNormal(poutine.block(model, hide=["x"])))
+    return guide
+
+
 @pytest.mark.parametrize("auto_class", [
     AutoDelta,
     AutoDiagonalNormal,
     AutoMultivariateNormal,
     AutoLowRankMultivariateNormal,
     auto_guide_list_x,
+    auto_guide_callable,
 ])
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
 def test_median(auto_class, Elbo):
@@ -204,6 +221,53 @@ def model():
     guide()
 
 
+@pytest.mark.parametrize("auto_class", [
+    AutoDelta,
+    AutoDiagonalNormal,
+    AutoMultivariateNormal,
+    AutoLowRankMultivariateNormal,
+])
+def test_callable(auto_class):
+
+    def model():
+        pyro.sample("x", dist.Normal(0., 1.))
+        pyro.sample("y", dist.MultivariateNormal(torch.zeros(5), torch.eye(5, 5)))
+
+    def guide_x():
+        x_loc = pyro.param("x_loc", torch.tensor(0.))
+        pyro.sample("x", dist.Delta(x_loc))
+
+    guide = AutoGuideList(model)
+    guide.add(guide_x)
+    guide.add(auto_class(poutine.block(model, expose=["y"]), prefix="auto_y"))
+    values = guide()
+    assert set(values) == set(["y"])
+
+
+@pytest.mark.parametrize("auto_class", [
+    AutoDelta,
+    AutoDiagonalNormal,
+    AutoMultivariateNormal,
+    AutoLowRankMultivariateNormal,
+])
+def test_callable_return_dict(auto_class):
+
+    def model():
+        pyro.sample("x", dist.Normal(0., 1.))
+        pyro.sample("y", dist.MultivariateNormal(torch.zeros(5), torch.eye(5, 5)))
+
+    def guide_x():
+        x_loc = pyro.param("x_loc", torch.tensor(0.))
+        x = pyro.sample("x", dist.Delta(x_loc))
+        return {"x": x}
+
+    guide = AutoGuideList(model)
+    guide.add(guide_x)
+    guide.add(auto_class(poutine.block(model, expose=["y"]), prefix="auto_y"))
+    values = guide()
+    assert set(values) == set(["x", "y"])
+
+
 def test_empty_model_error():
     def model():
         pass