TensorFlow：获取while循环内部渐变的方法吗？

Question

我想操纵while循环内部的渐变。实际目的是操纵循环神经网络行为，但这里是一个简化的例子来实现这个想法：

t_0 = tf.constant(0)
x_0 = tf.constant(1.0)
size = 3
x_ta = tf.TensorArray(tf.float32, size=size)

def cond(t, x_prev, x_ta):
    return tf.less(t, size)

def body(t, x_prev, x_ta):
    x = 2*x_prev
    x_ta = x_ta.write(t, x)
    return t + 1, x, x_ta

_, _, x_ta = tf.while_loop(cond, body, [t_0, x_0, x_ta])
x = x_ta.pack()

loss = x[size-1]

dloss_dx, = tf.gradients(loss, x)

with tf.Session() as sess:
    sess.run(tf.initialize_all_variables())
    print(sess.run(dloss_dx))

在此示例中，x[t] = 2*x[t-1]适用于所有t，因此我们喜欢 dloss_dx为[4., 2., 1.]。我们取而代之的是[ 0. 0. 1.]。我不认为这是一个错误，但认为这是一个相当意图的行为：我们实际上采用pack操作的衍生物，x[2]实际上并不依赖。< / p>

那么：有没有办法在这里获得[4., 2., 1.]？

注意：我想到的一个hacky解决方案是实际计算while循环内的渐变，然后计算累积产品（对应于链规则）。但是这个解决方案是有缺陷的：它不会推广到向量值输入，因为tf.gradients计算（加权）导数的总和（用于反向自动微分）而不是雅可比行列式。