使用tf.range在tf.function循环中操作列表报错，如何解决？

错误原因分析 #

• 示例1错误：Python原生列表不支持用Tensor类型作为索引，tf.range生成的i是Tensor对象，无法直接用来索引Python列表。
• 示例2错误：在tf.function的图模式下，循环体（while_body子图）内创建的张量被限制在子图作用域中，直接追加到Python列表会导致外层函数无法访问这些张量——Python列表不是TensorFlow的可追踪容器，无法跨子图传递张量。

正确实现方式 #

使用TensorFlow专门为图模式设计的tf.TensorArray替代Python列表，它支持动态写入指定索引和追加操作，完全兼容图模式的追踪机制，能实现和原列表一样的功能。

对应示例1（修改指定索引元素）

import tensorflow as tf

@tf.function
def my_function():
    # 初始化固定大小的TensorArray，指定元素数据类型和形状
    outputs = tf.TensorArray(dtype=tf.float32, size=10)
    for i in tf.range(10):
        # 写入指定索引位置，注意重新赋值返回的新TensorArray
        outputs = outputs.write(i, tf.expand_dims(tf.zeros([10, 10000]), axis=1))
    # 转换为与原需求一致的张量列表返回
    return outputs.stack().unstack()

对应示例2（追加元素）

import tensorflow as tf

@tf.function
def my_function():
    # 初始化动态大小的TensorArray，允许后续追加扩展
    outputs = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)
    for i in tf.range(10):
        # 按索引写入实现追加效果（索引从0开始递增）
        outputs = outputs.write(i, tf.expand_dims(tf.zeros([10, 10000]), axis=1))
    # 转换为张量列表返回
    return outputs.stack().unstack()

补充说明 #

• tf.TensorArray是不可变对象，write()操作会返回新的实例，必须重新赋值给变量。
• 如果最终需要的是高阶张量而非列表，可直接返回outputs.stack()，省去unstack()步骤。

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