TensorFlow 1.12：如何高效获取张量中非零最小元素的索引？

这个问题我之前在TF1.x的图模式下也碰到过，用Python的if-else在map_fn的lambda里确实行不通——因为图模式下的张量是符号化的，tf.not_equal返回的是布尔张量，不是可直接用于Python条件判断的布尔值。给你两个更优雅且兼容图模式（包括tf.Dataset.map场景）的解决方案：

方案一：用tf.where向量化替换0为极大值 #

这是最直接高效的方法，用向量化操作替代逐元素的map_fn，性能更好且完全符合图模式要求：

# 定义一个足够大的常量，注意要和tag_mask_sizes的 dtype 保持一致
large_value = tf.constant(9999999, dtype=tag_mask_sizes.dtype)
# 用tf.where将所有0替换为large_value，非零元素保持不变
tag_mask_sizes_suppressed = tf.where(
    tf.equal(tag_mask_sizes, 0),
    large_value,
    tag_mask_sizes
)
# 此时argmin会自动忽略被替换的0，找到非零元素中的最小值索引
smallest_mask_index = tf.argmin(tag_mask_sizes_suppressed)

为什么这个方法更好？tf.where是TensorFlow原生的向量化操作，比tf.map_fn的逐元素处理效率高得多，尤其是当张量维度较大时；同时它完全在图模式下运行，不会出现Python条件判断与张量不兼容的问题。

方案二：添加边界情况处理（可选） #

如果你的张量有可能全为0，可以额外加一个判断逻辑，避免返回错误的索引：

large_value = tf.constant(9999999, dtype=tag_mask_sizes.dtype)
# 检查张量中是否存在非零元素
has_non_zero = tf.reduce_any(tf.not_equal(tag_mask_sizes, 0))

# 用tf.cond分支处理两种情况
smallest_mask_index = tf.cond(
    has_non_zero,
    # 存在非零元素时，执行正常的索引查找
    lambda: tf.argmin(tf.where(tf.equal(tag_mask_sizes, 0), large_value, tag_mask_sizes)),
    # 全为0时，返回一个默认值（比如-1），你可以根据需求调整
    lambda: tf.constant(-1, dtype=tf.int64)
)

这个方案能让你的代码更健壮，避免在极端场景下出现不符合预期的结果。

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