Eager execution in tensorflow probability stops recording gradients at second iteration

Question

Somehow I don't get gradients in my train loop. Below is a simple example.

import tensorflow as tf
import tensorflow_probability as tfp
tf.enable_eager_execution()

hidden_size = 32
output_size = 1
m = tf.keras.Sequential(
[
            tfp.layers.DenseLocalReparameterization(hidden_size, tf.nn.leaky_relu),
            tfp.layers.DenseLocalReparameterization(hidden_size, tf.nn.leaky_relu),
            tfp.layers.DenseLocalReparameterization(output_size)
        ]
)

If I run a gradient recording step two times, the second times doesn't show any gradients. A list with None types is returned.

for _ in range(2):
    with tf.GradientTape() as tape:
        loss_value = m(tf.ones((1, 2))) * 2
        print(tape.gradient(loss_value, m.trainable_variables))

This is not the case if we replace the model m by a 'standard' tensorflow model, i.e.

m = tf.keras.Sequential([tf.keras.layers.Dense(1)])

I am using tensorflow=1.13.1 and tensorflow-probability=0.6.0

Answer 1

This seems to be a temporary bug. I uploaded to the nightly releases of tensorflow and tensorflow probability and the problem is solved.