Tensorflow 2的梯度计算在p处产生nan结果

Question

当x=0时，下面的简化代码为导数输出nan。我正在运行tensorflow 2.0.0。你知道吗

import tensorflow as tf

x = tf.Variable([[-1.0], [0.0], [1.0]])

with tf.GradientTape(persistent=True) as t:
    t.watch(x)
    # case 1: y = x^4
    # y = tf.reduce_sum(tf.pow(x, 4), axis=1) # gives nan for 2nd to 5th derivative at x=0
    # case 2: y = x + x^2 + x^3 + x^4
    y = tf.reduce_sum(tf.pow(x, [[1, 2, 3, 4]]), axis=1) # gives nan for 2nd to 5th derivative at x=0
    dy_dx = t.gradient(y, x)
    d2y_dx2 = t.gradient(dy_dx, x)
    d3y_dx3 = t.gradient(d2y_dx2, x)
    d4y_dx4 = t.gradient(d3y_dx3, x)
    d5y_dx5 = t.gradient(d4y_dx4, x)
del t

tf.print(y)
tf.print(tf.transpose(dy_dx)) # transpose only to fit on one line when printed
tf.print(tf.transpose(d2y_dx2))
tf.print(tf.transpose(d3y_dx3))
tf.print(tf.transpose(d4y_dx4))
tf.print(tf.transpose(d5y_dx5))

这将输出正确的值，除非x=0：

[0 0 4]
[[-2 1 10]]
[[8 -nan(ind) 20]]
[[-18 -nan(ind) 30]]
[[24 -nan(ind) 24]]
[[0 -nan(ind) 0]]

如果改为运行tf.pow(x, 4)情况，则nan只显示五阶导数：

[1 0 1]
[[-4 0 4]]
[[12 0 12]]
[[-24 0 24]]
[[24 24 24]]
[[-0 -nan(ind) 0]]

所以我的问题是：

• tensorflow文档没有明确说明pow函数支持两个大小不同的参数，但是第一个输出y是正确的。有人有这方面的经验吗？我希望所有3个输入x值的矩阵都提升到4次方。

• 梯度返回的nan值是我应该报告的bug吗？我确实发现了之前可能与此相关的问题，但已修复：https://github.com/tensorflow/tfjs/issues/346