如何在eager模式下遍历tf.Tensor？常见实现报错如何解决？

1. 列表推导式写法报错原因 #

你给函数加了@tf.function装饰器，该装饰器会将函数转换为静态计算图执行，静态图模式下不支持Python原生的迭代、列表推导式等操作遍历tf.Tensor类型，因此抛出OperatorNotAllowedInGraphError。即使全局开启eager模式，被@tf.function装饰的函数也会切换到图执行模式，无法直接用Python语法遍历Tensor。

2. tf.map_fn写法报错原因 #

dtype参数配置错误：你传入map_fn的是两个Tensor组成的元组，但你定义的lambda返回值是单个标量，而dtype填写了(tf.int64, tf.int64)，告知map_fn期望返回双元素的元组，两者嵌套结构不匹配，因此抛出结构不一致的报错。

额外提醒：你定义的indexs值为[1,2,3]，而probs每行的长度为3，索引范围是0~2，索引3会触发越界错误，需要根据实际需求调整索引值。

方案1：纯eager模式下直接遍历（不需要图执行加速） #

直接去掉@tf.function装饰器即可，eager模式下原生支持Python迭代语法遍历Tensor：

import tensorflow as tf
import numpy as np

probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([0, 1, 2])) # 调整为合法索引范围

def iterate_tensor(probs, indexs):
    return [output[label] for output, label in zip(probs, indexs)]

print(iterate_tensor(probs, indexs))

方案2：兼容图模式的map_fn写法 #

修正dtype参数为单个值即可：

@tf.function
def iterate_tensor(probs, indexs):
    return tf.map_fn(lambda i: i[0][i[1]], (probs, indexs), dtype=tf.int64)

result = iterate_tensor(probs, indexs)
print(result)

方案3：无遍历的向量化最优实现 #

你当前的需求是按行取对应索引的元素，完全可以用TensorFlow内置的批量索引操作实现，性能远高于循环遍历：

@tf.function
def iterate_tensor(probs, indexs):
    # 构造行索引+列索引的组合，批量取数
    return tf.gather_nd(probs, tf.stack([tf.range(tf.shape(probs)[0]), indexs], axis=1))

result = iterate_tensor(probs, indexs)
print(result)

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