def _get_traces(self, model, guide, args, kwargs):
"""
Runs the guide and runs the model against the guide with
the result packaged as a trace generator.
"""
if self.max_plate_nesting == float('inf'):
self._guess_max_plate_nesting(model, guide, args, kwargs)
if self.vectorize_particles:
guide = self._vectorized_num_particles(guide)
model = self._vectorized_num_particles(model)
# Enable parallel enumeration over the vectorized guide and model.
# The model allocates enumeration dimensions after (to the left of) the guide,
# accomplished by preserving the _ENUM_ALLOCATOR state after the guide call.
guide_enum = EnumMessenger(first_available_dim=-1 -
self.max_plate_nesting)
model_enum = EnumMessenger() # preserve _ENUM_ALLOCATOR state
guide = guide_enum(guide)
model = model_enum(model)
q = queue.LifoQueue()
guide = poutine.queue(guide,
q,
escape_fn=iter_discrete_escape,
extend_fn=iter_discrete_extend)
for i in range(1 if self.vectorize_particles else self.num_particles):
q.put(poutine.Trace())
while not q.empty():
yield self._get_trace(model, guide, args, kwargs)
def get_enum_traces(model, x):
guide_enum = EnumMessenger(first_available_dim=-2)
model_enum = EnumMessenger()
guide_ = guide_enum(
infer.config_enumerate(model.guide, "parallel", expand=True))
model_ = model_enum(model.model)
guide_trace = poutine.trace(guide_, graph_type="flat").get_trace(x)
model_trace = poutine.trace(pyro.poutine.replay(model_, trace=guide_trace),
graph_type="flat").get_trace(x)
return guide_trace, model_trace
def _sample_posterior(model, first_available_dim, temperature, *args,
**kwargs):
# For internal use by infer_discrete.
# Create an enumerated trace.
with poutine.block(), EnumMessenger(first_available_dim):
enum_trace = poutine.trace(model).get_trace(*args, **kwargs)
enum_trace = prune_subsample_sites(enum_trace)
enum_trace.compute_log_prob()
enum_trace.pack_tensors()
return _sample_posterior_from_trace(model, enum_trace, temperature, *args,
**kwargs)
