为何在tf.function中遍历张量长度返回占位符张量而非数值？

核心原因 #

tf.function会将Python代码转换为TensorFlow计算图执行（图模式），而普通Python代码默认是即时执行模式（Eager Mode），二者执行逻辑存在本质差异：

• 在Eager模式下（tf.function外部），tf.shape(S)[0]会立即计算出具体的Python整数（此处为2），range()基于该整数生成[0,1]序列，循环自然能打印出0和1。
• 在图模式下（tf.function内部），tf.shape(x)[0]是一个符号张量（仅代表计算图中的一个节点，无具体数值），而Python内置的range()仅能接收Python整数作为参数。此时range()无法解析符号张量的实际值，只能将其作为循环的符号化依据，导致循环内的print(i)打印的是张量节点本身，而非运行时的具体数值。

正确实现方式 #

要在tf.function中遍历张量长度，需使用TensorFlow原生循环API，或确保获取到静态形状的Python整数：

方式1：使用静态形状（适用于张量形状已知的场景） #

如果张量形状在定义时固定，可直接用.shape[0]获取静态形状（Python整数）：

@tf.function
def test(x):
    # x.shape[0]是静态形状，构建图时就能得到Python整数
    for i in range(x.shape[0]):
        print(i)

S = tf.random.uniform([2,2],0,1)
test(S)

输出：

0
1

方式2：使用TensorFlow动态循环API（适用于动态形状场景） #

如果张量形状是动态变化的（运行时才确定），要用tf.range配合tf.while_loop或tf.map_fn，同时使用tf.print（计算图原生打印操作）：

@tf.function
def test(x):
    # 生成动态长度的张量序列
    def loop_body(i):
        tf.print(i)
        return i + 1
    tf.while_loop(lambda i: i < tf.shape(x)[0], loop_body, [0])

S = tf.random.uniform([2,2],0,1)
test(S)

输出：

0
1

方式3：将动态形状转为Python整数（仅适用于形状可静态推断的场景） #

若需在图模式中获取动态形状的Python值，可使用tf.get_static_value：

@tf.function
def test(x):
    batch_size = tf.get_static_value(tf.shape(x)[0])
    for i in range(batch_size):
        print(i)

S = tf.random.uniform([2,2],0,1)
test(S)

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