Minibatch ADVI with Potential for Weighted Observations

In this example for using Minibatch with ADVI, it is emphasized that total_size needs to be set in order to properly calibrate the full likelihood. This is fine when the observable is a standard RandomVariable, which has observed and total_size parameters. However, it isn’t clear to me how this particular correction is to be done when using pm.Potential to weight observations.

For example, I have a gamma regression with weighted observations (weights W and observations Y), where I use

y = pm.Potential('y_logp_weighted', W*pm.Gamma.dist(mu=mu, sd=sigma).logp(Y))

Can anyone provide guidance for how to deal with this situation?

I’ve also tried using DensityDist, which does have an observed parameter, but fails in a strange way if one tries to pass a total_size argument

def weighted_gamma_logp(y, w):
    return w*pm.Gamma.dist(mu=mu, sd=sigma).logp(y)

y_logp = pm.DensityDist('y_logp_weighted', weighted_gamma_logp, observed={'y': Y_mb, 'w': W_mb}, total_size=Y.shape[0])

where Y_mb, W_mb are the Minibatches of the observed values Y and corresponding weights W. This results in a MissingInputError:

MissingInputError: Input 1 of the graph (indices start from 0), used to compute dot(Minibatch, alpha), was not provided and not given a value. Use the Theano flag exception_verbosity=‘high’, for more information on this error.

which goes away when removing the total_size argument.

Maybe a better approach is to weighted Y first and use the weighted_Y as observed? Is the weighting parameter also part of the free parameters that you want to model?