Using pyro.infer.Predictive with a model that utilizes the pyro.infer.reparam.SplitReparam reparameterizer raises an error as pyro.infer.Predictive tries to sample from the model in order to determine site shapes.
Straightforward sampling of a model that utilizes the pyro.infer.reparam.SplitReparam reparameterizer is not possible as this reparameterizer introduces sites with the pyro.distributions.ImproperUniform distribution, which does not support sampling.
Solution
Wrap the model with the pyro.poutine.InitMessenger effect handler during the site shapes determination phase. This solves the problem as the pyro.poutine.InitMessenger effect handler assigns values to the pyro.distributions.ImproperUniform sites before they are sampled.
This is a specific feature of pyro.poutine.ReparamMessenger, which applies initialization by pyro.poutine.InitMessenger before sampling, even if it appears last in the messenger stack (see #2876).
Problem
Using
pyro.infer.Predictivewith a model that utilizes thepyro.infer.reparam.SplitReparamreparameterizer raises an error aspyro.infer.Predictivetries to sample from the model in order to determine site shapes.Straightforward sampling of a model that utilizes the
pyro.infer.reparam.SplitReparamreparameterizer is not possible as this reparameterizer introduces sites with thepyro.distributions.ImproperUniformdistribution, which does not support sampling.Solution
Wrap the model with the
pyro.poutine.InitMessengereffect handler during the site shapes determination phase. This solves the problem as thepyro.poutine.InitMessengereffect handler assigns values to thepyro.distributions.ImproperUniformsites before they are sampled.This is a specific feature of
pyro.poutine.ReparamMessenger, which applies initialization bypyro.poutine.InitMessengerbefore sampling, even if it appears last in the messenger stack (see #2876).Testing
The fix can be verified by running