diff --git a/pyro/poutine/replay_messenger.py b/pyro/poutine/replay_messenger.py index 9c26490528..b54c5ce510 100644 --- a/pyro/poutine/replay_messenger.py +++ b/pyro/poutine/replay_messenger.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Dict, Optional from pyro.poutine.messenger import Messenger +from pyro.poutine.util import site_is_subsample if TYPE_CHECKING: import torch @@ -70,6 +71,10 @@ def _pyro_sample(self, msg: "Message") -> None: """ assert msg["name"] is not None name = msg["name"] + # Preserve manually supplied plate indices. Automatically generated + # subsamples start with no value and should still be replayed. + if site_is_subsample(msg) and msg["value"] is not None: + return None if self.trace is not None and name in self.trace: guide_msg = self.trace.nodes[name] if msg["is_observed"]: diff --git a/tests/infer/test_autoguide.py b/tests/infer/test_autoguide.py index 9b640e6fcb..925628bd1d 100644 --- a/tests/infer/test_autoguide.py +++ b/tests/infer/test_autoguide.py @@ -1058,6 +1058,24 @@ def model(): assert d.shape == sample_shape + (2, 3) +def test_replay_preserves_explicit_subsample(): + def model(subsample): + with pyro.plate("data", 6, dim=-1, subsample=subsample) as batch: + pyro.sample("z", dist.Normal(0, 1)) + return batch + + pyro.set_rng_seed(0) + subsample = torch.tensor([4, 0, 1]) + guide = AutoNormal(model) + guide_trace = poutine.trace(guide).get_trace(subsample) + model_trace = poutine.trace(poutine.replay(model, trace=guide_trace)).get_trace( + subsample + ) + + assert not torch.equal(guide_trace.nodes["data"]["value"], subsample) + assert_equal(model_trace.nodes["data"]["value"], subsample) + + @pytest.mark.parametrize( "auto_class", [