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:
@eb8680 suggested using this in [funsor](https://github.com/pyro-ppl/funsor).

@neerajprad to use this in [NumPyro](https://github.com/pyro-ppl/numpyro) you'll additionally need numpy versions of `torch.distributions.constraints` and `.transforms`. I believe you should be able to use `torch.distributions.constraint_registry`'s `biject_to` unmodified.

Resolving patch:
diff --git a/examples/minipyro.py b/examples/minipyro.py
--- a/examples/minipyro.py
+++ b/examples/minipyro.py
@@ -53,7 +53,8 @@ def guide(data):
 
         # 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
--- a/pyro/contrib/minipyro.py
+++ b/pyro/contrib/minipyro.py
@@ -14,6 +14,7 @@
 from __future__ import absolute_import, division, print_function
 
 from collections import OrderedDict
+import weakref
 
 import torch
 
@@ -25,7 +26,7 @@
 #     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 @@ def step(self, *args, **kwargs):
         # 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
--- /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)