Unexpected behavior of weight in KLDivergenceRegularizer

Setting weight=None or weight=1. in tfpl.KLDivergenceRegularizer gives an implied KL loss about double that of an explicit KL loss metric (tfd.kl_divergence) and also double that of a manually calculated KL loss metric. Setting weight=.5 makes all KL losses similar. How can I understand this weight behavior?

tf.version == 2.4.0
tfp.version == 0.11.0

Code to repoduce (based on the VAE example provided here https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Probabilistic_Layers_VAE.ipynb):

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp
tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions
tfkb = tf.keras.backend

# disable eager execution to use metrics functions defined below,
# while avoiding the issue mentioned here
# https://github.com/tensorflow/probability/issues/519
tf.compat.v1.disable_eager_execution()

# load dataset
datasets, datasets_info = tfds.load(name='mnist',
                                    with_info=True,
                                    as_supervised=False)

def _preprocess(sample):
  image = tf.cast(sample['image'], tf.float32) / 255.  # Scale to unit interval.
  image = image < tf.random.uniform(tf.shape(image))   # Randomly binarize.
  return image, image

train_dataset = (datasets['train']
                 .map(_preprocess)
                 .batch(256)
                 .prefetch(tf.data.experimental.AUTOTUNE)
                 .shuffle(int(10e3)))
eval_dataset = (datasets['test']
                .map(_preprocess)
                .batch(256)
                .prefetch(tf.data.experimental.AUTOTUNE))

# specify model
input_shape = datasets_info.features['image'].shape
encoded_size = 16
base_depth = 32

prior = tfd.Independent(tfd.Normal(loc=tf.zeros(encoded_size), scale=1),
                        reinterpreted_batch_ndims=1)

encoder = tfk.Sequential([
    tfkl.InputLayer(input_shape=input_shape),
    tfkl.Lambda(lambda x: tf.cast(x, tf.float32) - 0.5),
    tfkl.Conv2D(base_depth, 5, strides=1,
                padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2D(base_depth, 5, strides=2,
                padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2D(2 * base_depth, 5, strides=1,
                padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2D(2 * base_depth, 5, strides=2,
                padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2D(4 * encoded_size, 7, strides=1,
                padding='valid', activation=tf.nn.leaky_relu),
    tfkl.Flatten(),
    # use IndependentNormal instead of MultivariateNormalTriL for comparison 
    # with manual KL calculation    
    # tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(encoded_size),
    #            activation=None),
    # tfpl.MultivariateNormalTriL(
    #     encoded_size,
    #     activity_regularizer=tfpl.KLDivergenceRegularizer(prior)),
    tfkl.Dense(tfpl.IndependentNormal.params_size(encoded_size),
               activation=None),
    tfpl.IndependentNormal(
        encoded_size,
        activity_regularizer=tfpl.KLDivergenceRegularizer(prior, weight=None, 
                                                          use_exact_kl=True)),        
])

decoder = tfk.Sequential([
    tfkl.InputLayer(input_shape=[encoded_size]),
    tfkl.Reshape([1, 1, encoded_size]),
    tfkl.Conv2DTranspose(2 * base_depth, 7, strides=1,
                         padding='valid', activation=tf.nn.leaky_relu),
    tfkl.Conv2DTranspose(2 * base_depth, 5, strides=1,
                         padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2DTranspose(2 * base_depth, 5, strides=2,
                         padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2DTranspose(base_depth, 5, strides=1,
                         padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2DTranspose(base_depth, 5, strides=2,
                         padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2DTranspose(base_depth, 5, strides=1,
                         padding='same', activation=tf.nn.leaky_relu),
    tfkl.Conv2D(filters=1, kernel_size=5, strides=1,
                padding='same', activation=None),
    tfkl.Flatten(),
    tfpl.IndependentBernoulli(input_shape, tfd.Bernoulli.logits),
])

vae = tfk.Model(inputs=encoder.inputs,
                outputs=decoder(encoder.outputs[0]))

# define metrics functions  
def rec_loss(y_true, y_pred):
  return -vae.output.log_prob(y_true)
    
def kl_loss(y_true, y_pred):
  return tfd.kl_divergence(encoder.output, prior)

def kl_loss_manual(y_true, y_pred):
  # Kingma and Welling (2014), Auto-Encoding Variational Bayes, ICLR.
  # See Appendix B of https://arxiv.org/pdf/1312.6114.pdf
  mean = encoder.output.mean()
  var = encoder.output.variance()
  return -.5 * tfkb.sum(1. + tfkb.log(var) - tfkb.square(mean) - var, axis=1)

# train model
negloglik = lambda x, rv_x: -rv_x.log_prob(x)

vae.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-3),
            loss=negloglik,
            metrics=[rec_loss, kl_loss, kl_loss_manual])

_ = vae.fit(train_dataset,
            epochs=3,
            validation_data=eval_dataset)

# calculate ratio of implied KL over calculated KL
hist = vae.history.history
print('ratio:', (np.array(hist['loss']) - np.array(hist['rec_loss']))/ np.array(hist['kl_loss']))

Output:

Train on 235 steps, validate on 40 steps
Epoch 1/3
235/235 [==============================] - ETA: 0s - batch: 117.0000 - size: 1.0000 - loss: 189.9439 - rec_loss: 173.2091 - kl_loss: 8.3689 - kl_loss_manual: 8.3689

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:2325: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.
  warnings.warn('`Model.state_updates` will be removed in a future version. '

235/235 [==============================] - 14s 33ms/step - batch: 117.0000 - size: 1.0000 - loss: 189.9439 - rec_loss: 173.2091 - kl_loss: 8.3689 - kl_loss_manual: 8.3689 - val_loss: 134.8216 - val_rec_loss: 108.7525 - val_kl_loss: 13.0691 - val_kl_loss_manual: 13.0691
Epoch 2/3
235/235 [==============================] - 12s 32ms/step - batch: 117.0000 - size: 1.0000 - loss: 129.3110 - rec_loss: 99.8012 - kl_loss: 14.7685 - kl_loss_manual: 14.7685 - val_loss: 124.9163 - val_rec_loss: 94.1113 - val_kl_loss: 15.4014 - val_kl_loss_manual: 15.4014
Epoch 3/3
235/235 [==============================] - 12s 32ms/step - batch: 117.0000 - size: 1.0000 - loss: 123.1625 - rec_loss: 91.7423 - kl_loss: 15.7247 - kl_loss_manual: 15.7247 - val_loss: 120.7175 - val_rec_loss: 88.3529 - val_kl_loss: 16.1751 - val_kl_loss_manual: 16.1751
ratio: [1.99963829 1.99816107 1.99813944]

Start by running the supplied VAE reproduction from the linked Probabilistic_Layers_VAE.ipynb example and inspect tfp.layers.KLDivergenceRegularizer with weight=None, 1, and .5. Compare its implied loss with tfd.kl_divergence and the manual calculation. Done means the weight behavior and the observed factor of two are explained and verified.