Handler to replace `sample` sites with different distribution

[tilmantroester]

My (limited) understanding is that lift allows me to change param sites into sample sites, while substitute replaces sample sites with fixed values. Is there an effect handler that allows me to replace sample sites with other distributions? I thought this might be useful when exploring different priors for a model.

I had a quick stab at naively merging the lift and substitute handlers by adding

if isinstance(value, numpyro.distributions.Distribution):
    msg["fn"] = value

to substitute.process_message.

This seems to work but I don't know enough about numpyro to judge where this might break. lift also has some _samples_cache functionality that's missing here. Is there a reason substitute doesn't already support this?

[fehiepsi]

lift is helpful for bayesian neural network models. It might be better to add some logic to your model to control the behavior. Using effect-handler seems unnecessary to me. Something like

def model(a_dist):
    return sample("a", a_dist)

then you can change the distribution via argument, rather than using an effect handler.