Why does only DenseVariational provide the kl_weight parameter?

The DenseVariational provides the kl_weight parameter, whose documentation is

kl_weight: Amount by which to scale the KL divergence loss between prior and posterior.

In the paper "Weight Uncertainty in Neural Networks", equation 9

includes a scaling factor (pi) that weights the complexity cost (the KL term) relative to the likelihood cost on each mini-batch. The authors of the mentioned paper say that there are different ways of choosing this scaling factor pi. They show two ways: pi=1/M (uniform scaling) and pi=(2^{M-i})/(2^M - 1), where M is the number of mini-batches in an epoch and i is the ith mini-batch. Anyway, the scaling factor pi depends, in both cases, on the number of mini-batches.

The need for this scaling factor comes from the assumption that the KL term in the ELBO is calculated given all training instances, while, in mini-batch stochastic gradient descent, we do not use all training instances, but only of a subset of those. Therefore, given that the KL term is added to the loss at every step of gradient descent (so at every mini-batch), we need to scale the KL term proportional to the number of mini-batches. In the issue https://github.com/tensorflow/probability/issues/651, I wonder how we should scale the KL term, given that some people suggest to scale it by the total number of training instances.

Questions

1. Does the parameter kl_weight of DenseVariational correspond to pi in equation 9 of the paper mentioned above? In other words, is it the number that is multiplied by the KL term, at each training iteration (i.e. mini-batch)? If not, then what does it represent?

2. Why only DenseVariational provides a way of specifying kl_weight? Why don't other variational layers, such as tfp.layers.DenseReparameterization, tfp.layers.DenseLocalReparameterization, etc., also provide such a parameter? I know I can provide the function kernel_divergence_fn, so there's a possibility to also scale the KL divergence term in these cases, but why don't you provide the possibility to specify the parameter kl_weight, which I assume, in this second question, to correspond to pi in equation 9 of paper mentioned above?

Start with the DenseVariational, DenseReparameterization, and DenseLocalReparameterization API documentation and their kernel_divergence_fn entry points. Determine how kl_weight relates to the paper's pi and whether the other layers should expose an equivalent option. Done means documenting the behavior and resolving both questions consistently.