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

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Mac OS Catalina (Version: 10.15.2 (19C57))
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.1
• Python version: 3.7.5
• GPU model and memory: Intel Iris Pro 1536 MB

Describe the current behavior

I get the error

AttributeError: ‘Tensor’ object has no attribute ‘log_prob’

with TensorFlow Probability 0.9 (and TF 2.1).

Describe the expected behavior

No error.

Code to reproduce the issue

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()

Comments

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.

(This is a duplicate issue of https://github.com/tensorflow/probability/issues/742, but, for completeness, I decided to open it here too.)