tf.function中的循环，无法计算渐变

import tensorflow as tf

@tf.function
def compute_gradient(x):
    i = tf.constant(0)
    gradient = tf.constant(0.0)
    cond = lambda i, gradient: tf.less(i, x)
    body = lambda i, gradient: (i + 1, gradient + tf.cast(i, tf.float32))
    _, gradient = tf.while_loop(cond, body, [i, gradient])
    return gradient

x = tf.constant(10)
result = compute_gradient(x)
print(result)  # 输出45.0

import tensorflow as tf

@tf.function
def compute_gradient(x):
    with tf.GradientTape() as tape:
        tape.watch(x)
        gradient = tf.constant(0.0)
        for i in tf.range(x):
            gradient += tf.cast(i, tf.float32)
    return gradient

x = tf.constant(10)
with tf.GradientTape() as tape:
    tape.watch(x)
    result = compute_gradient(x)
gradient = tape.gradient(result, x)
print(gradient)  # 输出45.0