Applying an MCMC Step to Sample from Distribution with KL Divergence Term

Thanks for the tip about set_subtensor! Do I understand correctly that the inplace flag needs to be set to True in order to have an effect on \rm KL(q||p)

I experimented with a few more things (pushed to the same link as above) and I suspect I am now seeing the same shape error you are, having to do with the size of dist_params and ensuring it’s consistent with ndims_params and ndim_supp. My understanding from reading the docs for CustomDist/DensityDist was that

  • dist_params are extra arguments passed to logp, but are not themselves treated as random variables. In our case, this would be dist_params = [lam, nmc], where lam is \lambda in the paper and nmc is the number of monte carlo samples used to evaluate \rm KL(q||p). Or do I misunderstand the semantics of dist_params? Do the params include the RV? I noticed that the shape error that I’m seeing happens because pymc is trying to infer the shape of my RV using the dist_params, which makes me think I am not understanding what dist_params should contain.
  • I set ndims_params to None because the docs say that they will be assumed to be scalars, which they are (assuming the params are just (lam, nmc)).
  • I set ndim_supp to 1 because the variational parameters \theta=[\mu,\rho] can be treated as a 1D array. Or should this be ndim_supp=len(mu)+len(rho) indicating that the variational parameters live in \mathbb{R}^d with d=|\mu|+|\rho|?

tl;dr I suppose I remain confused about how exactly “parameters” and “dimensions” should be defined in CustomDist.