pyro-ppl / pyro

Deep universal probabilistic programming with Python and PyTorch

http://pyro.ai

Apache License 2.0

8.58k stars 987 forks source link

Fix models not rendering after application of the equalize effect handler #3387

Closed BenZickel closed 4 months ago

BenZickel commented 4 months ago

Problem

After using the pyro.poutine.equalize effect handler model rendering with pyro.render_model does not work due to no base_dist attribute in Delta sampled sites.

Solution

Add masking so that the sampling function has the base_dist attribute as expected by pyro.render_model for Delta sampled sites.

BenZickel commented 4 months ago

Also added a fix for this lint failure due to a mypy update.