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

The following code

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd


def get_mnist_data(normalize=True):
    img_rows, img_cols = 28, 28
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

    if tf.keras.backend.image_data_format() == 'channels_first':
        x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
        x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
        input_shape = (1, img_rows, img_cols)
    else:
        x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
        x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
        input_shape = (img_rows, img_cols, 1)

    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')

    if normalize:
        x_train /= 255
        x_test /= 255

    return x_train, y_train, x_test, y_test, input_shape


def get_bayesian_cnn(input_shape, num_classes=10):
    model_input = tf.keras.layers.Input(shape=input_shape)

    # kernel_divergence_fn=None to solve a symbolic exception.
    x = tfp.layers.Convolution2DFlipout(6, kernel_size=(5, 5), padding="SAME", activation=tf.nn.relu,
                                        kernel_divergence_fn=None)(model_input)
    x = tf.keras.layers.Flatten()(x)
    x = tfp.layers.DenseFlipout(84, activation=tf.nn.relu)(x)
    x = tfp.layers.DenseFlipout(num_classes)(x)

    model_output = tfp.layers.DistributionLambda(lambda t: tfd.Categorical(logits=t, validate_args=True))(x)

    model = tf.keras.Model(model_input, model_output)

    return model


def neg_log_likelihood(y_true, y_pred):
    return -tf.reduce_mean(y_pred.log_prob(tf.cast(tf.argmax(y_true, axis=-1), tf.int32)))


def train():
    x_train, y_train, x_test, y_test, input_shape = get_mnist_data()

    model = get_bayesian_cnn(input_shape=input_shape)

    model.compile(optimizer=tf.keras.optimizers.Adam(), loss=neg_log_likelihood,
                  metrics=[neg_log_likelihood])

    model.fit(x_train, y_train, batch_size=128, epochs=1, verbose=1)


if __name__ == "__main__":
    train()

produces the error

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

with TF 2.1 and TFP 0.9.

This error seems to be due to the fact that y_pred is a tensor when the loss is called, while it should be a distribution. Meanwhile, I found a question on Stack Overflow related to the third issue I mentioned above.

Start with the get_bayesian_cnn and neg_log_likelihood functions in the issue body, and reproduce the example with TensorFlow 2.1 and TensorFlow Probability 0.9. Trace how the model output is passed to the loss; done means the reported AttributeError is resolved and the example trains with the intended loss behavior.