Preserve explicit plate subsamples during replay

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"]:

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",
     [