Support data subsampling

[fehiepsi]

Addresses #556. This will also be helpful for #724.

@fritzo The implementation here is a bit simpler than Pyro. But I am not sure if I miss any point. Could you help me review this PR?

[eb8680]

Looks mostly correct to me, but you might want to add tests that exercise replay and nested/reused plates, like the model in this Pyro test.

The implementation here is a bit simpler than Pyro. But I am not sure if I miss any point.

Part of the extra complexity in Pyro comes from supporting sequential plates, which I guess isn't a goal here, and part of it comes from treating subsample sites as sample sites, which was just a poor early design choice in Pyro. Also, are you planning to add a numpyro.subsample primitive?

[fehiepsi]

add tests that exercise replay and nested/reused plates

Thanks @eb8680 ! I didn't even think about this. Thanks for pointing it out. :)

planning to add a numpyro.subsample primitive?

I didn't have that plan. I already took a look at its implementation but didn't understand how it works yet. Let me see if it is easy to address.

[eb8680]

Let me see if it is easy to address.

The implementation in Pyro is here in SubsampleMessenger._postprocess_message: https://github.com/pyro-ppl/pyro/blob/dev/pyro/poutine/subsample_messenger.py#L134

Here's a simplified version that does not include parameter subsampling:

    def postprocess_message(self, msg):  # new method for plate
        if msg["type"] == "subsample" 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:
                        statement = "pyro.subsample(..., event_dim={})".format(event_dim)
                        raise ValueError(
                            "Inside pyro.plate({}, {}, dim={}) invalid shape of {}: {}"
                            .format(self.name, self.size, self.dim, statement, shape))
                    if self.subsample_size < self.size:
                        value = msg["value"]
                        new_value = value.index_select(dim, self._indices)
                        msg["value"] = new_value

[fehiepsi]

Thanks for the snippet, Eli! My intention is just making a quick PR for this but it is worth to extend this to support subsample and param primitives. Also Pyro allows subsample_size is None in model and not None in guide. I remember @fritzo said that he will use NumPyro if we have this feature implemented. That's a good motivation. :))

[fehiepsi]

@eb8680 I have added subsample primitive and support subsample for it and param. Also made a small change to allow subsample_size=None in model but not None in guide.

[fritzo]

Note sure if this first subsample PR is the right place for thorough tests of SVI, but I found Pyro's test_subsample_gradient() to be invaluable in detecting subsampling issues with very little compute cost. In particular it helped us correctly scale the different parts of the ScoreParts tuple.

Does this PR aim to support subsampling in SVI? If so do you already have good tests for poutine.scale()?

[fehiepsi]

@fritzo Thanks for reviewing! I just add subsample_gradient test. We have some tests for handlers.scale in test_mask and test_scale, but we haven't supported ScoreParts stuff in NumPyro yet.

Does this PR aim to support subsampling in SVI?

My main motivation is to support #724, where subsample is required (see this comment). But SVI should work I believe.

[fritzo]

Looks great, thanks for implementing this!

[fritzo]

LGTM, @eb8680 did you have any other comments?

[eb8680]

Looks good, thanks for adding the extra tests!

[fehiepsi]

Thanks for reviewing, @eb8680 and @fritzo! I just observed something wrong in Predictive for models with multiple subsample statements Edit turns out to be my mistake, only replay handler needs to be fixed. I'll try to fix it and add the failing test in this PR.

[fehiepsi]

This test exercises subsample primitive works under replay handler.

[fehiepsi]

I switch the order of two sides of and because some primitives such as subsample do not have name field.