如何计算 tensorflow_probability 层上的梯度？

如何解决如何计算 tensorflow_probability 层上的梯度？

我想使用 tensorflow_probability 计算 tf.GradientTape() 层上的梯度。这是相当简单的使用正常的，例如，密集层

inp = tf.random.normal((2,5))
layer = tf.keras.layers.Dense(10)

with tf.GradientTape() as tape:
    out = layer(inp)
    loss = tf.reduce_mean(1-out)
grads = tape.gradient(loss,layer.trainable_variables)
print(grads)
[<tf.Tensor: shape=(5,10),dtype=float32,numpy=
 array([[ 0.04086879,0.04086879,-0.02974391,-0.07061271],[ 0.01167339,0.01167339,-0.02681615,-0.03848954],[ 0.00476769,0.00476769,-0.00492069,-0.00968838],[-0.00462376,-0.00462376,0.05914849,0.06377225],[-0.11682947,-0.11682947,-0.06357963,0.05324984]],dtype=float32)>,<tf.Tensor: shape=(10,),numpy=
 array([-0.05,-0.05,-0.1,-0.05],dtype=float32)>]

但是如果我使用 DenseReparameterization 来执行此操作，则 grads 不会注册。

inp = tf.random.normal((2,5))
layer = tfp.layers.DenseReparameterization(10)

with tf.GradientTape() as tape:
    out = layer(inp)
    loss = tf.reduce_mean(1-out)
grads = tape.gradient(loss,layer.trainable_variables)
print(grads)
[None,None,None]

谁能告诉我如何解决这个问题，以便对梯度进行录音和注册？

解决方法

啊哈，就是这样！我正在使用 tf v2.1.0。显然这不适用于 tensorflow_probability。我会尽快升级。谢谢 gobrewers14。