Using poutine.condition and subsample consistently

Pyro seems to encourage the use of effect handlers and, indeed, I generally prefer to use poutine.condition rather than the keyword argument obs in pyro.sample.
Typical use cases where this is fantastic are to quickly convert a generative model into a conditioned model for inference or to compare a model where a variable is assumed known vs one where it is latent.

However, poutine.condition does not seem to be compatible with subsampling (e.g. to train using mini-batches) and when I need to use mini-batch training I am forced to change framework and use obs.

To clarify what I mean, consider the code below, it does not return a valid model that can be used for inference.

def generative_model():
    feature_plate = pyro.plate("feature_plate", size=90, dim=-1)
    sample_plate = pyro.plate("sample_plate", size=1000, subsample_size=20, dim=-2)
    mu = pyro.sample("mu", pyro.distributions.Normal(0, 1))
    with feature_plate:
        with sample_plate:
            X = pyro.sample("X", pyro.distributions.Normal(mu, 1))

Xdata = pyro.distributions.Normal(0, 1).sample((1000,90))
inference_model = poutine.condition(generative_model, data={"X": Xdata})

Is there a way to use the effect handler logic to condition a model (i.e. poutine.condition) still allow mini-batch training (automatic or custom)?

I was thinking a way of doing it that would be consistent with pyro’s effect handler composability would be:

inference_model = poutine.condition(generative_model, data={"X": pyro.subsample(Xdata, 0)})

That of course does not work because the. pyro.subsample statement is not inside a pyro.plate.

Is there a solution right now or am I going in a “feature request” direction with this topic?

Hi @Gioelelm this is a great question and deserves a feature request issue on github. I agree poutine.condition is conceptually cleaner, though I often use obs= myself.

One way to view this is as a question about commutativity of effect handlers. We’ve had similar issues and done work to ensure InitMessenger commutes with ReparamMessenger. My guess is that we might be able to add logic to either SubsampleMessenger or ConditionMessenger to ensure the two commute.

:loudspeaker: Calling all programming language theorists, we’d love to learn about general theory of commutativity of effect handlers. If you are doing research in this area we’re happy to provide concrete real world examples!

I have added a github issue. Thanks!