I’m trying to implement the local reparameterization trick (Kingma, Salimans, and Welling 2015) to fit regressions with the spikeandslab prior (pointnormal mixture) on regression coefficients. I previously did this in Theano using the reparameterization gradient and analytical KL.
My current attempt in Edward is https://github.com/aksarkar/nwas/blob/4d6a1332eb39ca2b5876e14912cbf8eae1b2ed3f/analysis/example.org
Is there a better way to do this in Edward?
I defined two new Edward random variables:

SpikeSlab
, which supports analytical KL for the prior 
GeneticValue
(domainspecific jargon for x * theta), which supports sampling (using atf.contrib.distributions.Normal
instance) and dummy analytical KL (just returns 0)
Then, I call ed.ReparameterizationKLKLqp
directly since I know it won’t blow up.
I can’t think of a way to do this that doesn’t expose the reparameterization in the model specification, but this solution doesn’t play nicely with ed.copy
, so evaluating the model (e.g. computing coefficient of determination) requires pulling out the coefficients and computing things outside of Edward/Tensorflow.