TensorFlow中使用tf.function计算梯度返回None的问题咨询

• 核心原因是@tf.function编译生成的静态计算图对外部的tf.GradientTape为黑盒结构：当你在tape上下文内调用编译后的函数f(a)时，外部tape只能观测到输入是可训练变量a，输出是b、c两个张量，无法感知函数内部b依赖c计算的关联关系，梯度追踪链路断裂，最终返回None。
• 无装饰器的fplain运行在Eager即时执行模式下，所有操作逐行在tape上下文内执行，c的计算、b基于c的计算全程被tape追踪，梯度链路完整，因此可以得到预期结果。

如果需要在@tf.function场景下获取b对c的梯度，直接将梯度计算逻辑移到函数内部即可：

@tf.function
def f_with_grad(a):
    with tf.GradientTape() as internal_tape:
        c = a * 2
        b = tf.reduce_sum(c ** 2 + 2 * c)
    grad_b_c = internal_tape.gradient(b, c)
    return b, c, grad_b_c

a = tf.Variable([[0., 1.], [1., 0.]])
b, c, grad = f_with_grad(a)
print(grad) 
# 输出和Eager模式一致：<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
# array([[2., 6.],
#        [6., 2.]], dtype=float32)>

TensorFlow的@tf.function设计逻辑是优先保证静态图的执行效率，默认不会暴露内部张量的依赖关系给外部上下文，除非显式在函数内部完成梯度计算，这种设计可以避免不必要的梯度缓存占用内存。

内容的提问来源于stack exchange，提问作者marlon