DistributionLambda regularizer behavior

Hi,

I'm trying to use tfp.layers.DistributionLambda with a loss regularizer embedded in it, and am finding what appear to be strange results.
Taking a distribution generator:

def my_dist(inputs=None):
    if inputs is None:
        mean, var = tf.zeros(10), tf.ones(10)
    else:
        mean, var = tf.split(inputs, [10, 10], axis=-1)
    return tfd.MultivariateNormalDiag(loc=mean, scale_diag=var)

I first calculated a KL-divergence loss myself via:

>>> params = np.r_[np.zeros(10, dtype='f4') + 0.1, np.ones(10, dtype='f4')*1.1]
>>> print(tfd.kl_divergence(my_dist(params), my_dist()))
tf.Tensor(0.14689809, shape=(), dtype=float32)

which matches what I expect. Then, if I construct a layer and pass the params through:

layer = tfp.layers.DistributionLambda(
    make_distribution_fn=my_dist,
    activity_regularizer=tfp.layers.KLDivergenceRegularizer(
        my_dist(),
        use_exact_kl=True,
        weight=1,
    ),
)
layer(params)

I find the loss

>>> layer.losses
[<tf.Tensor: id=39326, shape=(), dtype=float32, numpy=-0.05400238>,
 <tf.Tensor: id=39420, shape=(), dtype=float32, numpy=-0.05400238>]

and it changes each time, so I assume use_exact_kl is being ignored and I try to approximate it:

>>> layer(tf.tile(params[None, ...], (10000, 1)))
>>> layer.losses
[<tf.Tensor: id=41627, shape=(), dtype=float32, numpy=0.14320771>,
 <tf.Tensor: id=41721, shape=(), dtype=float32, numpy=0.14320771>]

which seems to recover the expected value, but now there are two copies. Am I missing something in the way I construct the distribution?

Start by reproducing the provided DistributionLambda and KLDivergenceRegularizer snippets, comparing the scalar KL result with layer.losses for single and tiled inputs. Then inspect the DistributionLambda regularizer path and determine why the loss varies and appears twice; done means the behavior is explained and the regularizer produces the intended result or its documented limitation is clear.