pyro-ppl / numpyro / 42599031860

93%

push

github

fix(infer): use the guide's subsample indices in TraceEnum_ELBO (#2415)

The funsor trace handler only recorded sample and param sites, so the guide trace built by TraceEnum_ELBO had no plate sites. replay then had no subsample indices to hand to the model, and the model drew its own mini-batch. Local latents sampled by the guide for one set of indices were scored against data at different indices, which biases the loss and its gradients whenever the guide has sites inside a subsampled plate.

Record plate sites in the funsor trace, as numpyro.handlers.trace does, so the model reuses the guide's indices like it does in Trace_ELBO.

1 of 1 new or added line in 1 file covered. (100.0%)

15951 of 17157 relevant lines covered (92.97%)

1.7 hits per line