AttributeError: 'Tensor' object has no attribute 'log_prob' with ModelCheckpoint callback and save_best_only=True

Hi, I believe I found a bug while using TensorFlow probability layers and the ModelCheckpoint Keras callback with the option save_best_only=True.

I was able to reproduce this bug with the VAE example on Google Colab with both Tensorflow version 2.2 and tf-nightly builds. I also tried to update TensorFlow probability to version 0.10.1 but the bug remains.

From the VAE example on google colab you just need to add the ModelCheckpoint callback to the model.fit().

negloglik = lambda x, rv_x: -rv_x.log_prob(x)
vae.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-3),
            loss=negloglik)

callbacks = [tf.keras.callbacks.ModelCheckpoint('./checkpoint/', monitor='val_loss', save_best_only=True, verbose=0)]
_ = vae.fit(train_dataset,
            epochs=15,
            validation_data=eval_dataset,
            callbacks=callbacks)

Here is the reported error:

AttributeError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/distribution_layer.py:1275 __call__  *
        return self._kl_divergence_fn(distribution_a)
    /usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/distribution_layer.py:1394 _fn  **
        kl = kl_divergence_fn(distribution_a, distribution_b_)
    /usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/layers/distribution_layer.py:1378 kl_divergence_fn
        distribution_a.log_prob(z) - distribution_b.log_prob(z),

    AttributeError: 'Tensor' object has no attribute 'log_prob'

Note that if you set save_weights_only=True, save_best_only=True the code runs just fine.
Thanks,
Igor.

Start with the VAE example and reproduce the failure by adding the ModelCheckpoint callback to model.fit with save_best_only=True; compare it with the save_weights_only=True case. Inspect the TensorFlow Probability distribution_layer.py path shown in the traceback and verify that full checkpointing completes without the Tensor-versus-distribution AttributeError.