Keras model.compile uses incorrect loss when using Bayesian layers.

When I use my custom loss function, I got a wrong loss output if I choose keras.compile:

import tensorflow as tf

from tensorflow.examples.tutorials.mnist import input_data as mnist_data

train, valid, test = mnist_data.read_data_sets('~/code/Python')
num_classes = 10
from tensorflow import keras
import tensorflow_probability as tfp
model = keras.Sequential()

model.add(tfp.layers.DenseReparameterization(10, activation = 'softmax', input_shape=(784,)))

sgd = keras.optimizers.SGD(lr=.1, momentum=0.9, nesterov=True)
def my_loss(y_true,y_pred):
    return tf.keras.losses.categorical_crossentropy(y_true,y_pred)
model.compile(loss=my_loss, optimizer=sgd, metrics=['accuracy'])

x_train, y_train = train.images, train.labels
x_test, y_test = test.images, test.labels

y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
model.fit(x_train, y_train,
              batch_size=128,
              epochs=10,
              validation_data=(x_test, y_test),
              shuffle=True)

However, If I replace the bayesian layer by a traditional Dense layer, it seems everything is correct.