Separate query operations for imputing missing data

[eb8680]

Many real-world causal inference tasks, such as the mediation analysis example #14, involve working with datasets that have randomly or systematically missing data.

Pyro is already well-equipped to support simultaneous inference and imputation thanks to the obs_mask keyword argument to pyro.sample. That includes systematic missingness of the form discussed by Pearl in Graphical Models for Inference with Missing Data
because the value of obs_mask can depend on random variables and other computations within a (causal) Pyro program.

However, obs_mask cannot be read or modified by effect handlers or reparameterizers. IIRC this was a design choice made to ensure that imputation composed correctly with the rest of Pyro (closely related to the discussions in #30 and #82), but as a result it is currently somewhat difficult to separate the specification of a causal model, a set of observed values, and a pattern of missingness.

The semantics of obs_mask is quite simple, so to get around this limitation we could do something like the following:

1. Define a reparameterizer closely based on the current implementation of obs_mask handling that takes an obs_mask value for a single sample site as input
2. Define an effect handler or reparam.Strategy similar to the existing effect handler pyro.poutine.mask for higher-level specification of patterns of missingness across multiple sample sites

Some design questions:

• What information would the higher-level handler require from a user?
• What (non-?)commutativity laws should it obey wrt pyro.poutine.mask and MultiWorldCounterfactual? See Specify algebraic relationships between transformations #30
• Can we provide some degree of automation to users to ensure it is applied correctly? See Higher-level syntax for composing query operations #82

[eb8680]

What information would the higher-level handler require from a user?

To illustrate why this may be less than expected, consider the following interface which uses a single scalar placeholder value to infer an obs_mask value from observation tensors:

def missing(missval: float) -> ContextManager:
    ...

MISSVAL = 0.12345
x_obs[torch.isnan(x_obs)] = MISSVAL

conditioned_model = missing(missval=MISSVAL)(
    pyro.condition(data={"x": x_obs})(
        model))

# conditioned_model is equivalent to:
def model():
    ...
    x = pyro.sample("x", ..., obs=x_obs, obs_mask=(x_obs == MISSVAL))
    ...

Here is a sketch of an implementation of missing that reuses the existing obs_mask machinery in Pyro:

class missing(pyro.poutine.messenger.Messenger):

    def __init__(self, missval: float):
        self.missval = missval
        super().__init__()

    def _pyro_sample(self, msg):
        if (
            msg["is_observed"]
            and msg["value"] is not None
            and not pyro.poutine.util.site_is_subsample(msg)
            and not site_is_deterministic(msg)
        ):
            msg["stop"] = True
            msg["done"] = True
            with pyro.poutine.messenger.block_messengers(lambda m: m is self):
                msg["value"] = pyro.sample(
                    msg["name"],
                    msg["fn"],
                    obs=msg["value"],
                    obs_mask=(msg["value"] != self.missval),
                    infer=msg["infer"].copy(),
                )