Repository: pyro-ppl/pyro

Issue:
Support masked sample statement for batched imputation
Pyro can currently impute data in sample statements via
```py
def model(x=None):
    x = pyro.sample("x", my_dist, obs=x)
    ...
```
where `x` is either sampled when calling `model()` or observed when calling `model(observed_value)`.

It would be nice if this supported batching, where within a single batch some values are observed and some need to be sampled, say specified by `obs_mask`. Pyro can already handle masked observe statements via `TorchDistribution.mask()`, which allow tensorized decision between whether an observe statement exists, i.e. whether the `log(p)` and `log(q)` terms exist. What is needed is a different tensorized decision between whether a statement is a sample statement or an observe statement, which determines only whether the `log(q)` exists.

A possible implementation could be
```py
def sample(name, fn, obs=None, obs_mask=None):
    ...
    value = fn()
    if obs_mask is not None and obs_mask.any():
        value[obs_mask] = obs[obs_mask]        # use obs where available; sample otherwise
        guide_site["fn"] = fn.mask(~obs_mask)  # exclude log(q) terms at observed entries
    else:
        guide_site["fn"] = fn
    model_site["fn"] = fn
    ...
```


Discussion:
Note nan issues with backward through `dist.Categorical` as [mentioned on the forume](https://forum.pyro.ai/t/parent-nodes-with-missing-data/2004/4).

Resolving patch:
diff --git a/pyro/primitives.py b/pyro/primitives.py
--- a/pyro/primitives.py
+++ b/pyro/primitives.py
@@ -7,6 +7,8 @@
 from contextlib import ExitStack, contextmanager
 from inspect import isclass
 
+import torch
+
 import pyro.distributions as dist
 import pyro.infer as infer
 import pyro.poutine as poutine
@@ -61,6 +63,30 @@ def param(name, *args, **kwargs):
     return _param(name, *args, **kwargs)
 
 
+def _masked_observe(name, fn, obs, obs_mask, *args, **kwargs):
+    # Split into two auxiliary sample sites.
+    with poutine.mask(mask=obs_mask):
+        observed = sample(f"{name}_observed", fn, *args, **kwargs, obs=obs)
+    with poutine.mask(mask=~obs_mask):
+        unobserved = sample(f"{name}_unobserved", fn, *args, **kwargs)
+
+    # Interleave observed and unobserved events.
+    shape = obs_mask.shape + (1,) * fn.event_dim
+    batch_mask = obs_mask.reshape(shape)
+    try:
+        value = torch.where(batch_mask, observed, unobserved)
+    except RuntimeError as e:
+        if "must match the size of tensor" in str(e):
+            shape = torch.broadcast_shapes(observed.shape, unobserved.shape)
+            batch_shape = shape[:len(shape) - fn.event_dim]
+            raise ValueError(
+                f"Invalid obs_mask shape {tuple(obs_mask.shape)}; should be "
+                f"broadcastable to batch_shape = {tuple(batch_shape)}"
+            ) from e
+        raise
+    return deterministic(name, value)
+
+
 def sample(name, fn, *args, **kwargs):
     """
     Calls the stochastic function `fn` with additional side-effects depending
@@ -72,14 +98,25 @@ def sample(name, fn, *args, **kwargs):
     :param fn: distribution class or function
     :param obs: observed datum (optional; should only be used in context of
         inference) optionally specified in kwargs
+    :param ~torch.Tensor obs_mask: Optional boolean tensor mask of shape
+        broadcastable with ``fn.batch_shape``. If provided, events with
+        mask=True will be conditioned on ``obs`` and remaining events will be
+        imputed by sampling. This introduces a latent sample site named ``name
+        + "_unobserved"`` which should be used by guides.
+    :type obs_mask: bool or ~torch.Tensor
     :param dict infer: Optional dictionary of inference parameters specified
         in kwargs. See inference documentation for details.
     :returns: sample
     """
+    # Transform obs_mask into multiple sample statements.
     obs = kwargs.pop("obs", None)
-    infer = kwargs.pop("infer", {}).copy()
-    # check if stack is empty
+    obs_mask = kwargs.pop("obs_mask", None)
+    if obs_mask is not None:
+        return _masked_observe(name, fn, obs, obs_mask, *args, **kwargs)
+
+    # Check if stack is empty.
     # if stack empty, default behavior (defined here)
+    infer = kwargs.pop("infer", {}).copy()
     if not am_i_wrapped():
         if obs is not None and not infer.get("_deterministic"):
             warnings.warn("trying to observe a value outside of inference at " + name,


Test patch:
diff --git a/tests/infer/mcmc/test_valid_models.py b/tests/infer/mcmc/test_valid_models.py
--- a/tests/infer/mcmc/test_valid_models.py
+++ b/tests/infer/mcmc/test_valid_models.py
@@ -412,3 +412,28 @@ def potential_fn(z):
 
     mcmc.run()
     mcmc.get_samples()['points']
+
+
+@pytest.mark.parametrize("mask", [
+    torch.tensor(True),
+    torch.tensor(False),
+    torch.tensor([True]),
+    torch.tensor([False]),
+    torch.tensor([False, True, False]),
+])
+@pytest.mark.parametrize("Kernel, options", [
+    (HMC, {"adapt_step_size": True, "num_steps": 3}),
+    (NUTS, {"adapt_step_size": True, "max_tree_depth": 3}),
+])
+def test_obs_mask_ok(Kernel, options, mask):
+    data = torch.tensor([7., 7., 7.])
+
+    def model():
+        x = pyro.sample("x", dist.Normal(0., 1.))
+        with pyro.plate("plate", len(data)):
+            y = pyro.sample("y", dist.Normal(x, 1.),
+                            obs=data, obs_mask=mask)
+            assert ((y == data) == mask).all()
+
+    mcmc_kernel = Kernel(model, max_plate_nesting=0, **options)
+    assert_ok(mcmc_kernel)
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
@@ -2181,6 +2181,110 @@ def guide():
     assert_ok(model, guide, Elbo())
 
 
+@pytest.mark.parametrize("num_particles", [1, 2])
+@pytest.mark.parametrize("mask", [
+    torch.tensor(True),
+    torch.tensor(False),
+    torch.tensor([True]),
+    torch.tensor([False]),
+    torch.tensor([False, True, False]),
+])
+@pytest.mark.parametrize("Elbo", [
+    Trace_ELBO,
+    TraceEnum_ELBO,
+    TraceGraph_ELBO,
+    TraceMeanField_ELBO,
+])
+def test_obs_mask_ok(Elbo, mask, num_particles):
+    data = torch.tensor([7., 7., 7.])
+
+    def model():
+        x = pyro.sample("x", dist.Normal(0., 1.))
+        with pyro.plate("plate", len(data)):
+            y = pyro.sample("y", dist.Normal(x, 1.),
+                            obs=data, obs_mask=mask)
+            assert ((y == data) == mask).all()
+
+    def guide():
+        loc = pyro.param("loc", torch.zeros(()))
+        scale = pyro.param("scale", torch.ones(()),
+                           constraint=constraints.positive)
+        x = pyro.sample("x", dist.Normal(loc, scale))
+        with pyro.plate("plate", len(data)):
+            with poutine.mask(mask=~mask):
+                pyro.sample("y_unobserved", dist.Normal(x, 1.))
+
+    elbo = Elbo(num_particles=num_particles, vectorize_particles=True,
+                strict_enumeration_warning=False)
+    assert_ok(model, guide, elbo)
+
+
+@pytest.mark.parametrize("num_particles", [1, 2])
+@pytest.mark.parametrize("mask", [
+    torch.tensor(True),
+    torch.tensor(False),
+    torch.tensor([True]),
+    torch.tensor([False]),
+    torch.tensor([False, True, True, False]),
+])
+@pytest.mark.parametrize("Elbo", [
+    Trace_ELBO,
+    TraceEnum_ELBO,
+    TraceGraph_ELBO,
+    TraceMeanField_ELBO,
+])
+def test_obs_mask_multivariate_ok(Elbo, mask, num_particles):
+    data = torch.full((4, 3), 7.0)
+
+    def model():
+        x = pyro.sample("x", dist.MultivariateNormal(torch.zeros(3), torch.eye(3)))
+        with pyro.plate("plate", len(data)):
+            y = pyro.sample("y", dist.MultivariateNormal(x, torch.eye(3)),
+                            obs=data, obs_mask=mask)
+            assert ((y == data).all(-1) == mask).all()
+
+    def guide():
+        loc = pyro.param("loc", torch.zeros(3))
+        cov = pyro.param("cov", torch.eye(3),
+                         constraint=constraints.positive_definite)
+        x = pyro.sample("x", dist.MultivariateNormal(loc, cov))
+        with pyro.plate("plate", len(data)):
+            with poutine.mask(mask=~mask):
+                pyro.sample("y_unobserved", dist.MultivariateNormal(x, torch.eye(3)))
+
+    elbo = Elbo(num_particles=num_particles, vectorize_particles=True,
+                strict_enumeration_warning=False)
+    assert_ok(model, guide, elbo)
+
+
+@pytest.mark.parametrize("Elbo", [
+    Trace_ELBO,
+    TraceEnum_ELBO,
+    TraceGraph_ELBO,
+    TraceMeanField_ELBO,
+])
+def test_obs_mask_multivariate_error(Elbo):
+    data = torch.full((3, 2), 7.0)
+    # This mask is invalid because it includes event shape.
+    mask = torch.tensor([[False, False], [False, True], [True, False]])
+
+    def model():
+        x = pyro.sample("x", dist.MultivariateNormal(torch.zeros(2), torch.eye(2)))
+        with pyro.plate("plate", len(data)):
+            pyro.sample("y", dist.MultivariateNormal(x, torch.eye(2)),
+                        obs=data, obs_mask=mask)
+
+    def guide():
+        loc = pyro.param("loc", torch.zeros(2))
+        x = pyro.sample("x", dist.MultivariateNormal(loc, torch.eye(2)))
+        with pyro.plate("plate", len(data)):
+            with poutine.mask(mask=~mask):
+                pyro.sample("y_unobserved", dist.MultivariateNormal(x, torch.eye(2)))
+
+    elbo = Elbo(strict_enumeration_warning=False)
+    assert_error(model, guide, elbo, match="Invalid obs_mask shape")
+
+
 @pytest.mark.parametrize("scale", [1, 0.1, torch.tensor(0.5)])
 @pytest.mark.parametrize("Elbo", [
     Trace_ELBO,