Implementing ADVI

Hi everyone, I’ve implemented ADVI (i think) using ReparameterizationKLqp class, in order to approximate the posterior of a Gamma variable, Z.
To do so, I define the input latent variables dictionary as {Z : LogNormal} where the log Normal distribution is devised as a TransformedDistribuition instance.

It works perfectly fine!

I’ve used the LogNormal because if we apply the log transformation to Z and define the variational distribution of the transformed variable as a Normal, the resulting Z’s distribution is a LogNormal
However, both key and value have the same (constrainted) support as such no transformation is performed, so

  1. how is it working fine, if in ADVI the latent variables should be unconstrained?
  2. Am I implementing ADVI correctly?
  3. Or should I pass as input {log(Z): Normal}?