TensorFlow 2.0如何获取张量中与[0,0,0]相等的行对应的索引

你可以通过逐元素比对、行维度聚合判断、提取索引的组合逻辑实现该需求，完整可运行代码如下：

import tensorflow as tf

# 初始化输入张量
a = tf.constant([[0, 0, 0], [0, 0, 1], [0, 0, 0], [0, 0, 1]])
# 定义要匹配的目标行
target_row = tf.constant([0, 0, 0])

# 生成行匹配布尔掩码：某行所有元素和目标行相等时为True
row_match_mask = tf.reduce_all(tf.equal(a, target_row), axis=1)
# 提取符合条件的索引，去除多余维度后得到一维结果
result_indices = tf.squeeze(tf.where(row_match_mask), axis=1)

# 若需要转为Python列表，可执行下行注释代码
# result_indices = result_indices.numpy().tolist()

print(result_indices)
# 输出：tf.Tensor([0 2], shape=(2,), dtype=int64)

步骤说明 #

• tf.equal(a, target_row)：将输入张量的每个元素和目标行对应位置元素做等值比对，输出和输入张量同形状的布尔张量
• tf.reduce_all(..., axis=1)：沿行维度（索引为1的维度）做逻辑与聚合，只有整行所有元素比对结果都为True时，聚合结果才为True，最终得到长度等于输入张量行数的一维布尔掩码
• tf.where(row_match_mask)：提取所有掩码为True的位置索引，默认返回形状为(匹配行数, 1)的张量，通过tf.squeeze去除最后一个冗余维度，得到一维的索引结果

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