KLDivergenceAddLoss not returning a distribution

I am getting an error when I use the KLDivergenceAddLoss layer. My understanding is that this should add a loss and return the distribution at the previous layer in the network. However, in the example below the second function raises an error due to the fact that y is a Tensor following the inclusion of KLDivergenceAddLoss in the network.

I am using

• Python 3.7.7
• tensorflow 2.1.0
• tensorflow-probability 0.9.0

import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions
tfpl = tfp.layers
tfk = tf.keras
tfkl = tf.keras.layers

def make_model_no_kl():
    x = tfk.Input([5])
    y = tfk.Sequential([
        tfkl.Dense(2),
        tfpl.DistributionLambda(
            make_distribution_fn=lambda t: tfd.Categorical(logits=t)
        )
    ])(x)
    print(y.sample())
    return tfk.Model(inputs=x, outputs=y)


def make_model_with_kl():
    x = tfk.Input([5])
    y = tfk.Sequential([
        tfkl.Dense(2),
        tfpl.DistributionLambda(
            make_distribution_fn=lambda t: tfd.Categorical(logits=t)
        ),
        tfpl.KLDivergenceAddLoss(tfd.Categorical(logits=[1,2]))
    ])(x)
    print(y.sample())
    return tfk.Model(inputs=x, outputs=y)

m1 = make_model_no_kl()
m2 = make_model_with_kl()

2020-04-05 22:34:49.027766: W tensorflow/python/util/util.cc:319] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
2020-04-05 22:34:49.031974: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2020-04-05 22:34:49.043548: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fb292f93880 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-04-05 22:34:49.043567: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
Tensor("Reshape_2:0", shape=(None,), dtype=int32)
Traceback (most recent call last):
  File "simple.py", line 35, in <module>
    m2 = make_model_with_kl()
  File "simple.py", line 31, in make_model_with_kl
    print(y.sample())
AttributeError: 'Tensor' object has no attribute 'sample'

Start by running the make_model_with_kl reproducer in simple.py with the listed TensorFlow and TensorFlow Probability versions, then inspect KLDivergenceAddLoss and the surrounding distribution-layer behavior. Done means the layer adds its loss while the model output retains the distribution behavior demonstrated by make_model_no_kl, including a working y.sample() call.