tracking variables in multivariate normal distributions

Dear developers,
Recently, I am trying to write code for calculating MLE via TFP.
I found that TFP will not track the loc parameter of multivariate normal when using GradientTape

Here is an example code:

import tensorflow_probability as tfp
import tensorflow as tf
tfd = tfp.distributions

dims = 4
mvn_model = tfd.MultivariateNormalTriL(
    loc=tf.Variable(tf.zeros([dims], dtype=tf.float32), name="mu"),
    scale_tril=tfp.util.TransformedVariable(
        tf.eye(dims, dtype=tf.float32),
        tfp.bijectors.FillScaleTriL(),
        name="raw_scale_tril"))

The mu and raw_scale_tril are both learnable, which can be checked by

print(mvn_model.trainable_variables)

However, in the training process with the following code

x = mvn_model.sample([1000, 1])
optimizer = tf.optimizers.Adam(learning_rate=1.)
with tf.GradientTape() as tape:
    loss_value = -tf.reduce_mean(tf.reduce_sum(mvn_model.prob(x), axis = 1))
print(tape.watched_variables())
gradients = tape.gradient(loss_value, mvn_model.trainable_variables)
optimizer.apply_gradients(zip(gradients, mvn_model.trainable_variables))

I found that the mu is not watched by tape anymore.
As a result, the training process cannot be finished successfully.
Similar code works for univariate normal.
If developers could figure out whether it is a bug or what I miss, this will be really helpful.

Best,

Run the provided MultivariateNormalTriL example and inspect GradientTape.watched_variables(), mvn_model.trainable_variables, and the resulting gradients. Compare this behavior with the univariate normal case, then determine whether loc tracking is still broken or whether the issue needs a documented usage correction; done means the cause and expected training behavior are verified.