Repository: pyro-ppl/pyro

Issue:
[FR] Support constraints in minipyro
Currently constraints in Pyro are handled by a `ParamStoreDict` class. The `ParamStoreDict` interface was simplified in #1363, but the class remains complex to retain backwards compatibility. It would be nice to pull out the minimal logic needed to handle constraints.

We could do this in minipyro in two different ways:
1. Simplify the `ParamStoreDict` to be a simple dict whose values are `(unconstrained_value, constraint)` pairs, and move `transform_to(constraint)` logic into `pyro.param()`.
    Advantages: less code overall
2. Keep the logic in `ParamStoreDict` but simplify its interface.
    Advantages: easier for users to use the `PARAM_STORE` directly

Discussion:
fritzo: @neerajprad after this merges I will copy the changes to funsor.minipyro. At some point we should also port the changes to numpyro.minipyro.
fritzo: > how important is it to maintain backwards compatibility

I agree with @eb8680 that in the case of *.minipyro it is more important to be able to write simple implementations than to support all previous features. This version is still compatible with the use `[pyro.param(name) for name in pyro.get_param_store()]`, and I think it's reasonable to limit the informal minipyro spec to this narrow usage.
eb8680: > get_param_store() API

I suppose we could support this by constructing a new dictionary:
```
def get_param_store():
    return {name: pyro.param(name) for name in PYRO_PARAM_STORE.keys()}
```

Patch:
diff --git a/examples/minipyro.py b/examples/minipyro.py
index b5044465..ecddf18a 100644
--- a/examples/minipyro.py
+++ b/examples/minipyro.py
@@ -53,7 +53,8 @@ def main(args):
 
         # Report the final values of the variational parameters
         # in the guide after training.
-        for name, value in pyro.get_param_store().items():
+        for name in pyro.get_param_store():
+            value = pyro.param(name)
             print("{} = {}".format(name, value.detach().cpu().numpy()))
 
         # For this simple (conjugate) model we know the exact posterior. In
diff --git a/pyro/contrib/minipyro.py b/pyro/contrib/minipyro.py
index 367190fc..d60bf8d2 100644
--- a/pyro/contrib/minipyro.py
+++ b/pyro/contrib/minipyro.py
@@ -14,6 +14,7 @@ found at examples/minipyro.py.
 from __future__ import absolute_import, division, print_function
 
 from collections import OrderedDict
+import weakref
 
 import torch
 
@@ -25,7 +26,7 @@ import torch
 #     See http://docs.pyro.ai/en/0.3.1/parameters.html
 
 PYRO_STACK = []
-PARAM_STORE = {}
+PARAM_STORE = {}  # maps name -> (unconstrained_value, constraint)
 
 
 def get_param_store():
@@ -167,25 +168,37 @@ def sample(name, fn, obs=None):
     return msg["value"]
 
 
-# param is an effectful version of PARAM_STORE.setdefault
+# param is an effectful version of PARAM_STORE.setdefault that also handles constraints.
 # When any effect handlers are active, it constructs an initial message and calls apply_stack.
-def param(name, init_value=None):
-
-    def fn(init_value):
-        value = PARAM_STORE.setdefault(name, init_value)
-        value.requires_grad_()
-        return value
+def param(name, init_value=None, constraint=torch.distributions.constraints.real):
+
+    def fn(init_value, constraint):
+        if name in PARAM_STORE:
+            unconstrained_value, constraint = PARAM_STORE[name]
+        else:
+            # Initialize with a constrained value.
+            assert init_value is not None
+            with torch.no_grad():
+                constrained_value = init_value.detach()
+                unconstrained_value = torch.distributions.transform_to(constraint).inv(constrained_value)
+            unconstrained_value.requires_grad_()
+            PARAM_STORE[name] = unconstrained_value, constraint
+
+        # Transform from unconstrained space to constrained space.
+        constrained_value = torch.distributions.transform_to(constraint)(unconstrained_value)
+        constrained_value.unconstrained = weakref.ref(unconstrained_value)
+        return constrained_value
 
     # if there are no active Messengers, we just draw a sample and return it as expected:
     if not PYRO_STACK:
-        return fn(init_value)
+        return fn(init_value, constraint)
 
     # Otherwise, we initialize a message...
     initial_msg = {
         "type": "param",
         "name": name,
         "fn": fn,
-        "args": (init_value,),
+        "args": (init_value, constraint),
         "value": None,
     }
 
@@ -247,7 +260,8 @@ class SVI(object):
         # Differentiate the loss.
         loss.backward()
         # Grab all the parameters from the trace.
-        params = [site["value"] for site in param_capture.values()]
+        params = [site["value"].unconstrained()
+                  for site in param_capture.values()]
         # Take a step w.r.t. each parameter in params.
         self.optim(params)
         # Zero out the gradients so that they don't accumulate.


Test patch:
diff --git a/tests/contrib/test_minipyro.py b/tests/contrib/test_minipyro.py
new file mode 100644
index 00000000..9e99b768
--- /dev/null
+++ b/tests/contrib/test_minipyro.py
@@ -0,0 +1,192 @@
+from __future__ import absolute_import, division, print_function
+
+import warnings
+
+import pytest
+import torch
+from torch.distributions import constraints
+
+from pyro.generic import distributions as dist
+from pyro.generic import infer, optim, pyro, pyro_backend
+from tests.common import assert_close, xfail_param
+
+# This file tests a variety of model,guide pairs with valid and invalid structure.
+# See https://github.com/pyro-ppl/pyro/blob/0.3.1/tests/infer/test_valid_models.py
+
+
+def assert_ok(model, guide, elbo, *args, **kwargs):
+    """
+    Assert that inference works without warnings or errors.
+    """
+    pyro.get_param_store().clear()
+    adam = optim.Adam({"lr": 1e-6})
+    inference = infer.SVI(model, guide, adam, elbo)
+    for i in range(2):
+        inference.step(*args, **kwargs)
+
+
+def assert_error(model, guide, elbo, match=None):
+    """
+    Assert that inference fails with an error.
+    """
+    pyro.get_param_store().clear()
+    adam = optim.Adam({"lr": 1e-6})
+    inference = infer.SVI(model,  guide, adam, elbo)
+    with pytest.raises((NotImplementedError, UserWarning, KeyError, ValueError, RuntimeError),
+                       match=match):
+        inference.step()
+
+
+def assert_warning(model, guide, elbo):
+    """
+    Assert that inference works but with a warning.
+    """
+    pyro.get_param_store().clear()
+    adam = optim.Adam({"lr": 1e-6})
+    inference = infer.SVI(model, guide, adam, elbo)
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        inference.step()
+        assert len(w), 'No warnings were raised'
+        for warning in w:
+            print(warning)
+
+
+@pytest.mark.parametrize("backend", ["pyro", "minipyro"])
+def test_generate_data(backend):
+
+    def model(data=None):
+        loc = pyro.param("loc", torch.tensor(2.0))
+        scale = pyro.param("scale", torch.tensor(1.0))
+        x = pyro.sample("x", dist.Normal(loc, scale), obs=data)
+        return x
+
+    with pyro_backend(backend):
+        data = model().data
+        assert data.shape == ()
+
+
+@pytest.mark.parametrize("backend", ["pyro", "minipyro"])
+def test_generate_data_plate(backend):
+    num_points = 1000
+
+    def model(data=None):
+        loc = pyro.param("loc", torch.tensor(2.0))
+        scale = pyro.param("scale", torch.tensor(1.0))
+        with pyro.plate("data", 1000, dim=-1):
+            x = pyro.sample("x", dist.Normal(loc, scale), obs=data)
+        return x
+
+    with pyro_backend(backend):
+        data = model().data
+        assert data.shape == (num_points,)
+        mean = data.sum().item() / num_points
+        assert 1.9 <= mean <= 2.1
+
+
+@pytest.mark.parametrize("backend", ["pyro", "minipyro"])
+def test_nonempty_model_empty_guide_ok(backend):
+
+    def model(data):
+        loc = pyro.param("loc", torch.tensor(0.0))
+        pyro.sample("x", dist.Normal(loc, 1.), obs=data)
+
+    def guide(data):
+        pass
+
+    data = torch.tensor(2.)
+    with pyro_backend(backend):
+        elbo = infer.Trace_ELBO()
+        assert_ok(model, guide, elbo, data)
+
+
+@pytest.mark.parametrize("backend", ["pyro", "minipyro"])
+def test_plate_ok(backend):
+    data = torch.randn(10)
+
+    def model():
+        locs = pyro.param("locs", torch.tensor([0.2, 0.3, 0.5]))
+        p = torch.tensor([0.2, 0.3, 0.5])
+        with pyro.plate("plate", len(data), dim=-1):
+            x = pyro.sample("x", dist.Categorical(p))
+            pyro.sample("obs", dist.Normal(locs[x], 1.), obs=data)
+
+    def guide():
+        p = pyro.param("p", torch.tensor([0.5, 0.3, 0.2]))
+        with pyro.plate("plate", len(data), dim=-1):
+            pyro.sample("x", dist.Categorical(p))
+
+    with pyro_backend(backend):
+        elbo = infer.Trace_ELBO()
+        assert_ok(model, guide, elbo)
+
+
+@pytest.mark.parametrize("backend", ["pyro", "minipyro"])
+def test_nested_plate_plate_ok(backend):
+    data = torch.randn(2, 3)
+
+    def model():
+        loc = torch.tensor(3.0)
+        with pyro.plate("plate_outer", data.size(-1), dim=-1):
+            x = pyro.sample("x", dist.Normal(loc, 1.))
+            with pyro.plate("plate_inner", data.size(-2), dim=-2):
+                pyro.sample("y", dist.Normal(x, 1.), obs=data)
+
+    def guide():
+        loc = pyro.param("loc", torch.tensor(0.))
+        scale = pyro.param("scale", torch.tensor(1.))
+        with pyro.plate("plate_outer", data.size(-1), dim=-1):
+            pyro.sample("x", dist.Normal(loc, scale))
+
+    with pyro_backend(backend):
+        elbo = infer.Trace_ELBO()
+        assert_ok(model, guide, elbo)
+
+
+@pytest.mark.parametrize("backend", [
+    "pyro",
+    xfail_param("minipyro", reason="not implemented"),
+])
+def test_local_param_ok(backend):
+    data = torch.randn(10)
+
+    def model():
+        locs = pyro.param("locs", torch.tensor([-1., 0., 1.]))
+        with pyro.plate("plate", len(data), dim=-1):
+            x = pyro.sample("x", dist.Categorical(torch.ones(3) / 3))
+            pyro.sample("obs", dist.Normal(locs[x], 1.), obs=data)
+
+    def guide():
+        with pyro.plate("plate", len(data), dim=-1):
+            p = pyro.param("p", torch.ones(len(data), 3) / 3, event_dim=1)
+            pyro.sample("x", dist.Categorical(p))
+        return p
+
+    with pyro_backend(backend):
+        elbo = infer.Trace_ELBO()
+        assert_ok(model, guide, elbo)
+
+        # Check that pyro.param() can be called without init_value.
+        expected = guide()
+        actual = pyro.param("p")
+        assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("backend", ["pyro", "minipyro"])
+def test_constraints(backend):
+    data = torch.tensor(0.5)
+
+    def model():
+        locs = pyro.param("locs", torch.randn(3), constraint=constraints.real)
+        scales = pyro.param("scales", torch.randn(3).exp(), constraint=constraints.positive)
+        p = torch.tensor([0.5, 0.3, 0.2])
+        x = pyro.sample("x", dist.Categorical(p))
+        pyro.sample("obs", dist.Normal(locs[x], scales[x]), obs=data)
+
+    def guide():
+        q = pyro.param("q", torch.randn(3).exp(), constraint=constraints.simplex)
+        pyro.sample("x", dist.Categorical(q))
+
+    with pyro_backend(backend):
+        elbo = infer.Trace_ELBO()
+        assert_ok(model, guide, elbo)