Using Effect Handlers to Adjust Priors?

Hi all.

I’m playing around with the effect handlers module, specifically the lift, substitute, do, condition handlers, experimenting with using them to extend a “stereotypical” model into a more flexible family. Knowing that substitute lets us change a sample to a param, and that lift lets us do the opposite, why is it not possible to change a sample of some distribution into a sample of another? i.e. use effect handlers to change a parameter’s prior distribution?

As an example, if I make a bare-bones simple model:

def model(X):
    x_mu  = numpyro.sample( 'x_mu',  dist.Uniform(-5.0,5.0))
    x_sig = numpyro.sample( 'x_sig', dist.Uniform( 0.0,5.0))

    with numpyro.plate('data',len(X)) :
        numpyro.sample('x', dist.Normal(x_mu,x_sig), obs=X)

And I want to change the prior of x_mu to something else, for example \mathcal{N}(-100,10), how can I do this? I’ve tried using lift, which would let me do this for a param:

lifted_model = lift(lifted_model, 
                    {'x_mu': dist.Normal(-100,10)}

I’ve also tried using substitute to convert it into a param first:

lifted_model = condition(model, 
                    {'x_mu': 10.0}

lifted_model = lift(lifted_model, 
                    {'x_mu': dist.Normal(-100,10)}

But in both cases the lift seems to be having no effect on the priors of lifted_model. Have I missed something here? There’s the obvious workaround of having the prior distribution as a model argument, but I’m hoping to avoid this unwieldy approach.

I might be a few months behind on my updates, so apologies if this has been changed in the last while.


This is not the case. substitute provides a value to a sample statement (instead of drawing a sample from prior_dist.sample(...)). If you want to replace prior, you can create an effect handler which changes msg['fn'] to something else.