TensorFlow 2.0如何按行相等条件从已有张量提取值生成新张量

报错原因 #

你代码中的range(h)是错误触发点：range仅接收整数参数，而h是张量类型，无法直接传入range生成循环序列，因此抛出类型错误。

正确实现方案 #

方案1：修改你的现有写法 #

调整循环范围参数，最后将收集的行拼接为张量即可：

import tensorflow as tf

# 输入张量定义
h = tf.constant([[0,0,0],[0,0,1],[0,1,0]])
h2 = tf.constant([[0,0,0],[0,0,1],[0,0,1]])

# 取张量的第一维长度（行数）作为循环范围
tensor_list = [h[i] for i in range(h.shape[0]) if tf.reduce_all(tf.equal(h[i], h2[i]))]
# 拼接列表得到最终张量
h3 = tf.stack(tensor_list)

方案2：推荐向量化实现（效率更高） #

直接使用TensorFlow原生的张量操作，无需Python循环，处理大张量时性能优势明显：

# 逐行判断是否完全相等，得到布尔掩码
equal_mask = tf.reduce_all(tf.equal(h, h2), axis=1)
# 用掩码过滤出相等的行
h3 = tf.boolean_mask(h, equal_mask)

两种方案输出的h3都符合预期结果：

tf.Tensor(
[[0 0 0]
 [0 0 1]], shape=(2, 3), dtype=int32)

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