[bug] AutoGuide silently overrides a model's manual `pyro.plate(subsample=idx)`

Issue Description

When a model manually builds its own subsample index and passes it explicitly via pyro.plate(name, size, dim=dim, subsample=idx), composing that model with a default AutoGuide (e.g. AutoNormal(model), no create_plates=) silently breaks the alignment between the model's and guide's minibatch.

Root cause, as far as I can trace it:

• AutoGuide._create_plates() (in pyro/infer/autoguide/guides.py) builds its own plate, when create_plates= is not supplied, for every plate name found in the model's prototype trace via:
  python复制

  self.plates[name] = pyro.plate(name, full_size, dim=frame.dim, subsample_size=frame.size)

Specifically, it uses subsample_size=, not subsample=. This makes the guide draw a fresh, independent random index on every call, with no reference to whatever idx the model itself was given.

• SVI/Trace_ELBO runs the guide first, then evaluates the model via poutine.replay(model, trace=guide_trace).
• ReplayMessenger._pyro_sample (pyro/poutine/replay_messenger.py) unconditionally does msg["value"] = guide_msg["value"] for any site name present in the guide's trace — including the plate's own internal subsample site (named after the plate, e.g. "data") — regardless of whether the model already supplied its own explicit value for that site.
• Net effect: the model's plate ends up silently using the guide's independently-drawn index instead of the idx the model was explicitly passed. Any plain Python indexing the model does with its own idx variable (e.g. data[idx]) is not affected by this (it's not routed through Pyro's effect handler stack), so the model's observations and the guide's/model's latent-variable row selection become decoupled — silently, with no exception, no warning, and a training run that looks completely normal.

I confirmed that passing create_plates= to the AutoGuide, reusing the model's exact plate (same subsample=idx), resolves the issue — the model's "data" site then correctly keeps the caller's idx. This matches the "useful for data subsampling" note already present in AutoGuide's docstring for create_plates=, but nothing about the default behavior warns that manual subsample=idx composed with a default AutoGuide is actively unsafe rather than merely "less convenient."

Environment

• OS: Linux-7.0.0-28-generic-x86_64-with-glibc2.39 (Ubuntu-based)
• Python: 3.10.16
• PyTorch: 2.13.0+cu130
• Pyro version: 1.9.1
• Relevant pip freeze:
  复制

  pyro-api==0.1.2
  pyro-ppl==1.9.1
  torch==2.13.0

Code Snippet

import torch
import pyro
import pyro.distributions as dist
import pyro.poutine as poutine
from pyro.infer.autoguide import AutoNormal

pyro.set_rng_seed(123)

N = 6
idx = torch.tensor([4, 0, 1])  # the caller's own, explicit, manually-built subsample


def model(idx):
    with pyro.plate("data", N, dim=-1, subsample=idx):
        pyro.sample("mu", dist.Normal(0.0, 1.0))


pyro.clear_param_store()
guide = AutoNormal(model)
guide(idx)  # sets up the guide's prototype trace

# This is exactly what SVI.step() / Trace_ELBO does internally:
#   1. trace the guide
#   2. run poutine.replay(model, trace=guide_trace)
guide_trace = poutine.trace(guide).get_trace(idx)
model_trace = poutine.trace(poutine.replay(model, trace=guide_trace)).get_trace(idx)

print("idx passed explicitly to model(idx):      ", idx.tolist())
print("guide's independently-drawn 'data' value: ", guide_trace.nodes["data"]["value"].tolist())
print("model's 'data' value after replay:        ", model_trace.nodes["data"]["value"].tolist())
assert torch.equal(model_trace.nodes["data"]["value"], guide_trace.nodes["data"]["value"])
assert not torch.equal(model_trace.nodes["data"]["value"], idx)
print("\nNo exception raised. The model's plate silently used the guide's")
print("unrelated random subsample instead of the idx it was explicitly given.")

Sample output (seed-dependent, but always shows the same divergence):

idx passed explicitly to model(idx):       [4, 0, 1]
guide's independently-drawn 'data' value:  [3, 4, 2]
model's 'data' value after replay:         [3, 4, 2]

No exception raised. The model's plate silently used the guide's
unrelated random subsample instead of the idx it was explicitly given.

Workaround: pass create_plates= to the AutoGuide, rebuilding the identical plate (same subsample=idx) so guide and model agree:

def create_plates(idx):
    return pyro.plate("data", N, dim=-1, subsample=idx)

guide = AutoNormal(model, create_plates=create_plates)