Equalize effect handler

[BenZickel]

• Implementation of Add effect handler that forces sample sites to sample the same value [feature request] #3374.
• See docs here and excerpt below.

Given a stochastic function with some primitive statements and a list of names, the equalize effect handler forces the primitive statements at those names to have the same value, with that value being the result of the first primitive statement matching those names.

Consider the following Pyro program:

def per_category_model(category):
    shift = pyro.param(f'{category}_shift', torch.randn(1))
    mean = pyro.sample(f'{category}_mean', pyro.distributions.Normal(0, 1))
    std = pyro.sample(f'{category}_std', pyro.distributions.LogNormal(0, 1))
    return pyro.sample(f'{category}_values', pyro.distributions.Normal(mean + shift, std))

Running the program for multiple categories can be done by

def model(categories):
    return {category:per_category_model(category) for category in categories}

To make the std sample sites have the same value, we can write

equal_std_model = pyro.poutine.equalize(model, '.+_std')

If on top of the above we would like to make the ‘shift’ parameters identical, we can write

equal_std_param_model = pyro.poutine.equalize(equal_std_model, '.+_shift', 'param')

[fritzo]

Thanks for the clean implementation.

From a software perspective this seems sensible, but from a modeling perspective I'm a bit worried that priors are double counted. I guess it's fine to trust the user.

[BenZickel]

Hi @fritzo, regarding your comment on double counting of priors, I think this is not happening, as only the first sample site is sampled from its distribution, while the remaining sites are converted to deterministic sample sites. Hopefully, this is what you meant in your comment.

[fritzo]

Ah I see, so the prior for each equalized site will be the prior from first matching site. Thanks!