Transformed Variable not trainable in Keras model

Hi,
I am trying to train a positive variable using tfp.util.TransformedVariable as an attribute of a tf.keras.Model object. However, the model does not recognize it as a trainable variable, and it does not receive gradients. This behavior holds for tensorflow_probability==0.10.0 and tensorflow==2.2.0, as well as for the nightly builds of both.

Here is a colab notebook illustrating this behavior: https://colab.research.google.com/drive/1XGCcm8l0OGRiy35lr3XcHAyZMuBNpsIB?usp=sharing

In this example, we are trying to train both an unconstrained variable (loc) and a constrained variable (scale). Only the loc variable updates.

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

class Model(tf.keras.Model):
  def __init__(self):
    super(Model, self).__init__()
    self.loc = tf.Variable(tf.ones(shape=[5]), name="loc")
    self.scale = tfp.util.TransformedVariable(
        tf.ones([5]),
        bijector=tfp.bijectors.Softplus(),
        name="scale") 
    self.distribution = tfp.distributions.Normal(loc=self.loc, scale=self.scale)

  def call(self, inputs):
    samples = self.distribution.sample()
    assigned_means = tf.gather(samples, inputs)
    return tfp.distributions.Normal(loc=assigned_means, scale=1.)

model = Model()
print(model.trainable_weights)  # only 'loc' shows up

optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)
loss = lambda x, rv: -tf.reduce_sum(rv.log_prob(x))
inputs = np.array([0, 1, 2, 3, 4]).astype(np.int32)
outputs = np.array([0., 1., 2., 3., 4.]).astype(np.float32)
dataset = tf.data.Dataset.from_tensor_slices((inputs, outputs))
dataset = dataset.batch(5)
model.compile(optimizer=optimizer, loss=loss)
model.fit(dataset, epochs=100, verbose=0)

# Check if the location parameters have moved from their original values.
assert(not (np.isclose(model.loc.numpy(), np.ones(5))).all())  # Passes

# Check if the scale parameters have moved from their original values.
assert(not (np.isclose(model.scale.numpy(), np.ones(5))).all())  # Fails

Thanks!

Start with the linked Colab notebook and reproduce the behavior using TensorFlow Probability 0.10.0 with TensorFlow 2.2.0 or the nightly builds. Inspect how TransformedVariable is tracked by the Keras model, comparing model.trainable_weights and gradients for loc and scale. Done means the constrained scale variable receives gradients and updates during fit, alongside loc.