Is there an idiomatic way to use baselines with reparameterizable variables? (For example, Reducing reparameterization gradient variance.) Or should I subclass TraceGraph_ELBO and implement this?
TraceGraph_ELBO appears to only use baselines if they are defined on non-reparameterizable variables. If baselines are provided on reparameterizable variables, they are ignored. (At least, that is my understanding based on the documentation and source code.)
Is something like what is described in the paper referenced already implemented under the hood in TraceGraph_ELBO? If not, is it because these methods not (anecdotally) helpful? Would an as-general-as-possible implementation of the method in the paper be a useful contribution to the pyro package?
Thanks