Incorrect training when using keras-style loss (with dist.log_prob)

I'm finding that my models are not training under specific, but easily reproducible circumstances. Doesn't result in runtime errors, but the models are not trained properly. Running Tensorflow 1.14 and tensorflow_probability 0.7.0
Here is an example:

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
from tensorflow_probability import bijectors as tfb
from tensorflow import keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras.datasets.mnist import load_data
import numpy as np

def get_loss(dist):
    def loss (y_true,y_pred):
        neg_log_likelihood = -tf.reduce_mean(dist.log_prob(y_true))
        return neg_log_likelihood
    return loss

sess= K.get_session()

((x_train,y_train),(x_test,y_test)) = load_data()
x_train = (x_train-127.5)/127.5
x_test = (x_test-127.5)/127.5
data_shape = x_train.shape[1:]
low = np.min(y_train)
high= np.max(y_train)
input_layer = keras.layers.Input(shape=data_shape)
x = keras.layers.Flatten()(input_layer)
x = keras.layers.Dense(512,activation='relu')(x)
x = keras.layers.Dense(256,activation='relu')(x)
x_out = keras.layers.Dense(high+1,activation='softmax')(x)

model = keras.models.Model([input_layer],[x_out])
model_output = model.output

dist = tfd.Categorical(probs=model_output)

model.compile(
            optimizer=keras.optimizers.Adam(),
            loss=get_loss(dist)
              )  
model.fit(x=x_train,y=y_train,validation_data=(x_test,y_test),
          batch_size=512,epochs=100,])

After some debugging I'm fairly sure that the issue lies specifically with how we set up the loss keras-style (the get_loss function).

if I compute the value of the tensor (instead of wrapping it up and feeding it to model.compile)
neg_log_likelihood = -tf.reduce_mean(dist.log_prob(y_true))
then obtained values are fine, the gradients seem fine too.

Any of the following things will make it run fine too:

• Training directly by using train_op = tf.compat.v1.train.AdamOptimizer().minimize(neg_log_likelihood)
• using loss=keras.losses.sparse_categorical_crossentropy instead of loss=get_loss(dist)
• Using dist = tfd.OneHotCategorical(probs=model_output) instead of dist = tfd.Categorical(probs=model_output)(after one hot encoding the targets)
• Similarly to above, using keras.losses.categorical_cross_entropy

Start with the supplied TensorFlow 1.14 and TensorFlow Probability 0.7.0 reproduction, then compare the Keras-style Categorical.log_prob loss with the listed working alternatives. No repository file or test is named; done means identifying and correcting the discrepancy so the model trains correctly and its loss and gradients agree with the direct training cases.