Higher-level syntax for composing query operations

[eb8680]

The design of this library emphasizes composing many different model transformations (effect handlers) like condition and do to generate a single expanded program representing a causal quantity of interest, but the syntax for applying a deeply nested sequence of handlers in Pyro is awkward and easy to mess up. The ITE query from the SLC example illustrates this:

def queried_model(data, *model_args, **model_kwargs):
        
    treatment_shape = (model.n_instances, model.n_objects)

    extended_model = do(actions=dict(T=torch.zeros(treatment_shape)))(
                        do(actions=dict(T=torch.ones(treatment_shape)))(
                            condition(data=data)(
                                model)))

    with MultiWorldCounterfactual(dim=model.outer_dim-1):
        Y_cf = extended_model(*model_args, **model_kwargs)
        
    return pyro.deterministic("ITEs", Y_cf[..., 1, 0, :, :] - Y_cf[..., 0, 1, :, :])

Coming up with general machinery for safely composing handlers or other program transformations is out of scope, certainly for first release.

Fortunately, to handle cases like the above it seems more straightforward to define a higher-order effect handler Query for composing our causal query transformations like do and MultiWorldCounterfactual as well as a few important ones built into Pyro such as condition and plate:

def Query(*handlers: pyro.poutine.messenger.Messenger) -> contextlib.ContextDecorator:
    ...
...
ite_query = Query(
    MultiWorldCounterfactual(dim=model.outer_dim-1),
    do(actions=dict(T=torch.zeros(treatment_shape))),
    do(actions=dict(T=torch.ones(treatment_shape))),
    condition(data=data),
)

queried_model = ite_query(model)

The following sketch of an implementation is almost sufficient to address the SLC example (except the final pyro.deterministic statement computing the ITE, discussed in #80, and the added argument data):

@contextlib.contextmanager
def Query(*handlers: pyro.poutine.messenger.Messenger):
    try:
        with contextlib.ExitStack() as stack:
            for h in reversed(handlers):
                stack.enter_context(h)
            yield stack
    finally:
        pass

This design raises several interesting questions:

• Should Query be able to consume other Query instances as arguments?
• Should Query be nestable inside models? If not, how should we detect and prevent this?
• Which upstream Pyro handlers could be allowed as arguments to Query without overly complicating its semantics? For example, what about reparam and block?

It also suggests a number of possible extensions or refinements:

• Could we reorganize or augment the type hierarchy of Query-compatible query handlers so that there is a common base type that can be subclassed by users?
• Could we extend the interface above to allow deferring large arguments like data as in the original example's queried_model?
• Could Query be a torch.nn.Module or pyro.nn.PyroModule for parameter and name compatibility, as part of Make Pyro more compatible with PyTorch idioms #44 and Add a "modern Pyro" programming tutorial #45 ?
• Could we also use Query as a place to address Specify algebraic relationships between transformations #30 by allowing it to fuse and canonically order query operations?
• Could we define the interface of Query in a way that enables model-agnostic specification of query templates like the ITE?
• Could we expose a Query-compatible interface to the select_world operation proposed in Computing causal quantities across worlds is too clunky #80?