@tf.function中遍历tf.Tensor报OperatorNotAllowedInGraphError的解决方案咨询

你遇到的报错本质是图模式下不支持使用Python原生列表推导、循环遍历动态Tensor，必须使用TensorFlow原生的张量操作实现同等筛选逻辑，以下是可直接运行的替代方案：

核心实现逻辑 #

不需要手动遍历索引，通过布尔掩码+向量化操作实现行筛选，兼容图模式且性能远高于逐行遍历：

1. 生成逐行匹配的布尔掩码：判断h和h2的每一行是否完全相等
2. 用掩码筛选h中符合条件的行
3. 若需对筛选后的行应用自定义函数fn，使用tf.map_fn（兼容图模式的映射接口）实现

完整示例代码 #

import tensorflow as tf

# 示例输入：形状均为[3,3]的张量
h = tf.constant([[1,2,3],[4,5,6],[7,8,9]])
h2 = tf.constant([[1,2,3],[0,0,0],[7,8,9]])

def fn(h_row):
    # 图模式下不要用Python原生print，改用tf.print
    tf.print(h_row)
    return h_row

@tf.function
def filter_matched_rows(h, h2):
    # 按行判断h和h2是否完全相等，生成布尔掩码
    match_mask = tf.reduce_all(tf.equal(h, h2), axis=1)
    # 筛选出匹配的行
    filtered_rows = tf.boolean_mask(h, match_mask)
    # 对每个匹配行应用fn函数，fn_output_signature要和fn返回值的形状、类型对应
    result = tf.map_fn(fn, filtered_rows, fn_output_signature=tf.TensorSpec(shape=(3,), dtype=h.dtype))
    return result

tensor = filter_matched_rows(h, h2)
tf.print(tensor)

备选实现（按索引筛选） #

如果你需要拿到匹配行的索引做额外处理，也可以用tf.where+tf.gather实现，效果和上述方案一致：

@tf.function
def filter_matched_rows(h, h2):
    match_mask = tf.reduce_all(tf.equal(h, h2), axis=1)
    # 拿到所有匹配行的索引
    matched_indices = tf.where(match_mask)[:, 0]
    # 按索引取对应行
    filtered_rows = tf.gather(h, matched_indices)
    result = tf.map_fn(fn, filtered_rows, fn_output_signature=tf.TensorSpec(shape=(3,), dtype=h.dtype))
    return result

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