Minibatch ADVI with Potential for Weighted Observations

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.