Replacing DenseVariational with DenseFlipout does not work

Hi,
I am trying to replace the DenseVariational layer with DenseFlipout in the Probabilistic Layers Regression Notebook.

Though, this causes model(x_tst).mean().numpy() to fail: AttributeError: 'Tensor' object has no attribute 'numpy'.

• With DenseVariational or keras.layers.Dense, everything works fine.
• With DenseReparameterization, DenseLocalReparameterization and DenseFlipout, I get the AttributeError

Am I missing something?

I am using the nightly versions:

tf version: 2.0.0-dev20190405
tfp version: 0.7.0-dev

Here is a minimal example (to run in the Colab):

xs = np.arange(15, dtype=np.float32).reshape((-1, 1))

model = tf.keras.Sequential([
    tfp.layers.DenseFlipout(1),
    tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1))
])
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=negloglik)
model.fit(xs, xs, epochs=100, verbose=False)

ys = model(xs)
assert isinstance(ys, tfd.Distribution)
print(ys.mean().numpy())  # AttributeError: 'Tensor' object has no attribute 'numpy'