Repository: pyro-ppl/pyro

Issue:
Params inside iarange
the current behavior of params created inside an `iarange` (aka vectorized map_data) is a bit nonsensical: a vector of params of length `batch_size` will be created and used on every iteration. this means that a param gets reused for whatever data point happens to be selected at that index in a batch. this comes up eg when there is a mean-field model for a local variable.

(there is a reasonable work-around, by creating the vector of needed params globally, and only selecting from it in the iarange.)

a better behavior would be to automatically associate params to their index into the data tensor, pulling out the appropriate params for a given batch (and creating them as needed). it might be tricky to get this working, and especially to do it efficiently.

i'm flagging this as a low priority bug for now. we can think through a solution after launch.


Discussion:
the other attitude to take is #alwaysamortize ; )
It seems pretty simple to write this as:
```py
with pyro.iarange("data", size, subsample_size) as ind:
    ps = pyro.param("ps", Variable(torch.ones(size), requires_grad=True)).index_select(ind)
    # ...do stuff with ps...
```
I believe the only way to get this to work automatically is to add a notion of `.event_dim` to parameters. However the `batch_dim`/`event_dim` distinction has proven to be confusing in the context of distributions. It seems like it will be easier for us to provide documentation for this subtle use of `param` inside `iarange`, rather than implementing automatic machinery that is itself confusing and needing documentation.

I'll remove this from the 0.2 milestone (if you don't mind). This will give us a chance to see how users adapt to other automatic dimension handling in e.g. enumeration. Given what we learn after 0.2 release, there may be a clear path to resolving this issue in Pyro 0.3.
We can wait until after 0.2, but I do think this is a significant issue, both because it makes irange semantically diverge from iarange and because it messes up mean field guides. 
Here is a proposal for refactoring towards a more plate-friendly `pyro.param` statement.

I think we could modify `param` behavior by inspecting the `cond_indep_stack`. For each vectorized frame we could `.expand(frame.size)` at the appropriate dim. Determining the appropriate dim is tricky because `param` currently does not have a notion of `event_dim`, and we would need to add something akin to `.independent()` or an `event_dim` kwarg to `param` statements. For each non-vectorized frame we could do some simple autoname stuff like appending `'_{}'.format(frame.counter)`. For consistency, we should probably also apply this autonaming to `sample` statements and avoid most of our current `'_{}'.format(i)` grossness.

To correctly handle subsampling, we would need to give `pyro.param` access to the subsampling tensors. This information can't live in the `cond_indep_stack`, but it could be looked-up in a trace via  `trace.nodes['_subsample_{}'.format(frame.name)]`. This is a little gross, but `param` could simply find that trace in the `_PYRO_STACK`. Once this information is available, vectorized param would slice its tensors by the subsample index, and sequential param would index into `subsample_index[frame.counter]` to create the correct param name.

To account for global params, we could use various autonaming hacks. E.g. we could avoid expanding in `vectorize_particles` if the `cond_indep_stack`'s `frame.name.startswith('_pyro_')`. Or we could add a kwarg `pyro.param(..., global=True)` to avoid the proposed `.expand()` behavior. @eb8680 could probably find cleaner patterns based on pyro.contrib.autoname.

@eb8680 @neerajprad @ngoodman WDYT?
what if we add an abstract index argument to param (and sample)? ie something along the lines of `pyro.param("foo", index=[I1 I2], ...)` where Ij can be either literals or abstract indices introduced by plates (assuming we go with the proposal to have plate bind a name for the index it introduces).

that is, within an iteration the indices are just strings/numbers that get appended to the name. within a plate block the index is an object that tells us which dimension to expand along. if a param is created without and index it is treated as global.

this by itself doesn't tell us what to do with subsampling, though as @eb8680 pointed out we should probably rethink that anyhow.... maybe @fritzo's suggestions is an answer?
I like the idea of a more plate friendly param!

>  This information can't live in the cond_indep_stack, but it could be looked-up in a trace via trace.nodes['_subsample_{}'.format(frame.name)]. This is a little gross, but param could simply find that trace in the _PYRO_STACK. 

It should mostly be fine, but I hope that the cost doesn't add up as we do this every time the model is run. In any case, I think we can explore other ideas of exposing this information to `param`.

> E.g. we could avoid expanding in vectorize_particles if the cond_indep_stack's frame.name.startswith('_pyro_'). Or we could add a kwarg pyro.param(..., global=True) to avoid the proposed 

Given our existing use cases, it might be easier to add a kwarg to `plate`, say `global_params=True`to denote that any params inside it do not need to be expanded. Then it will be easy for us to wrap a user model inside a `plate` without making any changes to the model code.

I have not thought through this, but I think this could also be extended to support iparam (e.g. by wrapping inside a plate with `global_params=False`) #1213, #1330.
Ok, I've illustrated possible syntax changes in #1491 (based on #1486 ). The intent is to:
1. illustrate current syntax in #1486;
2. work towards consensus on proposed syntax changes in #1491;
3. decide which changes we actually want to implement.
In addition to a `batch_shape`/`event_shape` distinction, it seems to me that there are at least two types of broadcasting operations we'd need to distinguish between - one that simply broadcasts and one that actually replicates the parameter tensor along a dimension to allow different values at each site.  This distinction corresponds to the distinction between `torch.Tensor.expand()` and `torch.Tensor.repeat()`.

For example, consider the case of SVI with vectorized `num_particles` and local parameters created by a `param` statement within a model `plate`.  Presumably we'd want to `repeat` along the model `plate`, but only `expand` along the `num_particles` `plate`.  It's easy enough to see that here, but what about more complicated settings, especially when a model may be used under multiple different broadcasting interpretations?

There's fundamentally ambiguity in the identity of each broadcasting dimension, and moreover `expand`/`repeat` decisions made in one setting will affect all future decisions that could be made about that parameter.  Is there a design decision we could make that would remove that ambiguity?  We can't simply create all parameters globally, since the existence, shape, or initial values of local parameters may depend on computation within a model.  Should we only allow `repeat` when a parameter is first created, and assume all broadcasts thereafter are `expand`s?  That's a possibility, but it introduces a dependence on the context in which the parameter was created, which sounds like a recipe for very tricky bugs.

Resolving patch:
diff --git a/pyro/params/param_store.py b/pyro/params/param_store.py
--- a/pyro/params/param_store.py
+++ b/pyro/params/param_store.py
@@ -183,7 +183,7 @@ def replace_param(self, param_name, new_param, old_param):
         assert self._params[param_name] is old_param.unconstrained()
         self[param_name] = new_param
 
-    def get_param(self, name, init_tensor=None, constraint=constraints.real):
+    def get_param(self, name, init_tensor=None, constraint=constraints.real, event_dim=None):
         """
         Get parameter from its name. If it does not yet exist in the
         ParamStore, it will be created and stored.
@@ -195,6 +195,7 @@ def get_param(self, name, init_tensor=None, constraint=constraints.real):
         :type init_tensor: torch.Tensor
         :param constraint: torch constraint
         :type constraint: torch.distributions.constraints.Constraint
+        :param int event_dim: (ignored)
         :returns: parameter
         :rtype: torch.Tensor
         """
diff --git a/pyro/poutine/subsample_messenger.py b/pyro/poutine/subsample_messenger.py
--- a/pyro/poutine/subsample_messenger.py
+++ b/pyro/poutine/subsample_messenger.py
@@ -3,6 +3,7 @@
 import torch
 
 from pyro.distributions.distribution import Distribution
+from pyro.poutine.util import is_validation_enabled
 from pyro.util import ignore_jit_warnings, jit_compatible_arange
 
 from .indep_messenger import CondIndepStackFrame, IndepMessenger
@@ -117,7 +118,7 @@ def _subsample(name, size=None, subsample_size=None, subsample=None, use_cuda=No
         return size, subsample_size, subsample
 
     def _reset(self):
-        self.subsample = None
+        self._indices = None
         super(SubsampleMessenger, self)._reset()
 
     def _process_message(self, msg):
@@ -128,3 +129,20 @@ def _process_message(self, msg):
                 with ignore_jit_warnings():
                     msg["scale"] = torch.tensor(msg["scale"])
         msg["scale"] = msg["scale"] * self.size / self.subsample_size
+
+    def _postprocess_message(self, msg):
+        if msg["type"] == "param" and self.dim is not None:
+            event_dim = msg["kwargs"].get("event_dim")
+            if event_dim is not None:
+                assert event_dim >= 0
+                dim = self.dim - event_dim
+                shape = msg["value"].shape
+                if len(shape) >= -dim and shape[dim] != 1:
+                    if is_validation_enabled() and shape[dim] != self.size:
+                        raise ValueError(
+                            "Inside pyro.plate({}, {}, dim={}) "
+                            "invalid shape of pyro.param({}, ..., event_dim={}): {}"
+                            .format(self.name, self.size, self.dim, msg["name"], event_dim, shape))
+                    # Subsample parameters with known batch semantics.
+                    if self.subsample_size < self.size:
+                        msg["value"] = msg["value"].index_select(dim, self._indices)
diff --git a/pyro/primitives.py b/pyro/primitives.py
--- a/pyro/primitives.py
+++ b/pyro/primitives.py
@@ -39,8 +39,22 @@ def param(name, *args, **kwargs):
     To interact with the param store or write to disk,
     see `Parameters <parameters.html>`_.
 
-    :param name: name of parameter
+    :param str name: name of parameter
+    :param init_tensor: initial tensor or lazy callable that returns a tensor.
+        For large tensors, it may be cheaper to write e.g.
+        ``lambda: torch.randn(100000)``, which will only be evaluated on the
+        initial statement.
+    :type init_tensor: torch.Tensor or callable
+    :param constraint: torch constraint, defaults to ``constraints.real``.
+    :type constraint: torch.distributions.constraints.Constraint
+    :param int event_dim: (optional) number of rightmost dimensions unrelated
+        to baching. Dimension to the left of this will be considered batch
+        dimensions; if the param statement is inside a subsampled plate, then
+        corresponding batch dimensions of the parameter will be correspondingly
+        subsampled. If unspecified, all dimensions will be considered event
+        dims and no subsampling will be performed.
     :returns: parameter
+    :rtype: torch.Tensor
     """
     kwargs["name"] = name
     return _param(name, *args, **kwargs)


Test patch:
diff --git a/tests/infer/test_gradient.py b/tests/infer/test_gradient.py
--- a/tests/infer/test_gradient.py
+++ b/tests/infer/test_gradient.py
@@ -48,11 +48,10 @@ def model(subsample):
             pyro.sample("x", Normal(z, 1), obs=x)
 
     def guide(subsample):
-        loc = pyro.param("loc", lambda: torch.zeros(len(data), requires_grad=True))
-        scale = pyro.param("scale", lambda: torch.tensor([1.0], requires_grad=True))
-        with pyro.plate("data", len(data), subsample_size, subsample) as ind:
-            loc_ind = loc[ind]
-            pyro.sample("z", Normal(loc_ind, scale))
+        scale = pyro.param("scale", lambda: torch.tensor([1.0]))
+        with pyro.plate("data", len(data), subsample_size, subsample):
+            loc = pyro.param("loc", lambda: torch.zeros(len(data)), event_dim=0)
+            pyro.sample("z", Normal(loc, scale))
 
     if scale != 1.0:
         model = poutine.scale(model, scale=scale)
diff --git a/tests/infer/test_valid_models.py b/tests/infer/test_valid_models.py
--- a/tests/infer/test_valid_models.py
+++ b/tests/infer/test_valid_models.py
@@ -215,6 +215,62 @@ def guide():
     assert_ok(model, guide, Elbo())
 
 
+@pytest.mark.parametrize("subsample_size", [None, 5], ids=["full", "subsample"])
+@pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
+def test_plate_subsample_param_ok(subsample_size, Elbo):
+
+    def model():
+        p = torch.tensor(0.5)
+        with pyro.plate("plate", 10, subsample_size):
+            pyro.sample("x", dist.Bernoulli(p))
+
+    def guide():
+        with pyro.plate("plate", 10, subsample_size) as ind:
+            p0 = pyro.param("p0", torch.tensor(0.), event_dim=0)
+            assert p0.shape == ()
+            p = pyro.param("p", 0.5 * torch.ones(10), event_dim=0)
+            assert len(p) == len(ind)
+            pyro.sample("x", dist.Bernoulli(p))
+
+    if Elbo is TraceEnum_ELBO:
+        guide = config_enumerate(guide)
+
+    assert_ok(model, guide, Elbo())
+
+
+@pytest.mark.parametrize("subsample_size", [None, 5], ids=["full", "subsample"])
+@pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
+@pytest.mark.parametrize("shape,ok", [
+    ((), True),
+    ((1,), True),
+    ((10,), True),
+    ((3, 1), True),
+    ((3, 10), True),
+    ((5), False),
+    ((3, 5), False),
+])
+def test_plate_param_size_mismatch_error(subsample_size, Elbo, shape, ok):
+
+    def model():
+        p = torch.tensor(0.5)
+        with pyro.plate("plate", 10, subsample_size):
+            pyro.sample("x", dist.Bernoulli(p))
+
+    def guide():
+        with pyro.plate("plate", 10, subsample_size):
+            pyro.param("p0", torch.ones(shape), event_dim=0)
+            p = pyro.param("p", torch.ones(10), event_dim=0)
+            pyro.sample("x", dist.Bernoulli(p))
+
+    if Elbo is TraceEnum_ELBO:
+        guide = config_enumerate(guide)
+
+    if ok:
+        assert_ok(model, guide, Elbo())
+    else:
+        assert_error(model, guide, Elbo(), match="invalid shape of pyro.param")
+
+
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
 def test_plate_no_size_ok(Elbo):
 
@@ -478,6 +534,34 @@ def model():
     assert_error(model, guide, Elbo(), match='hape mismatch inside plate')
 
 
+@pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
+def test_nested_plate_plate_subsample_param_ok(Elbo):
+
+    def model():
+        with pyro.plate("plate_outer", 10, 5):
+            pyro.sample("x", dist.Bernoulli(0.2))
+            with pyro.plate("plate_inner", 11, 6):
+                pyro.sample("y", dist.Bernoulli(0.2))
+
+    def guide():
+        p0 = pyro.param("p0", 0.5 * torch.ones(4, 5), event_dim=2)
+        assert p0.shape == (4, 5)
+        with pyro.plate("plate_outer", 10, 5):
+            p1 = pyro.param("p1", 0.5 * torch.ones(10, 3), event_dim=1)
+            assert p1.shape == (5, 3)
+            px = pyro.param("px", 0.5 * torch.ones(10), event_dim=0)
+            assert px.shape == (5,)
+            pyro.sample("x", dist.Bernoulli(px))
+            with pyro.plate("plate_inner", 11, 6):
+                py = pyro.param("py", 0.5 * torch.ones(11, 10), event_dim=0)
+                assert py.shape == (6, 5)
+                pyro.sample("y", dist.Bernoulli(py))
+
+    if Elbo is TraceEnum_ELBO:
+        guide = config_enumerate(guide)
+    assert_ok(model, guide, Elbo())
+
+
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
 def test_nonnested_plate_plate_ok(Elbo):