如何在TensorFlow中无需tf.map一次性比较不同形状张量？

现有两个不同形状的TensorFlow张量：

>>> A = tf.Tensor([4135047. 1193752.], shape=(2,), dtype=float32)
>>> B = tf.Tensor(
          [1226019. 4135047. 4135047. 4169911. 1193752. 4135047. 4135047. 4135047.],
          shape=(8,), dtype=float32
       )

当前的实现方式是分多次比较后合并结果：

>>> compare_1 = tf.math.equal(B, A[0])
tf.Tensor([False  True  True False False  True  True  True], shape=(8,), dtype=bool)

>>> compare_2 = tf.math.equal(B, A[1])
tf.Tensor([False False False False  True False False False], shape=(8,), dtype=bool)

# 最终结果
>>> tf.math.logical_or(compare_1, compare_2)
tf.Tensor([False  True  True False  True  True  True  True], shape=(8,), dtype=bool)

需求是不使用tf.map()，一次性完成比较：判断B的每个元素是否与A中的任意元素匹配，返回如下布尔张量：

>>> compare(B, A)
tf.Tensor([False  True  True False  True  True  True  True])

逻辑说明：

• B的第1个元素1226019与A中任意元素都不匹配 → False
• B的第2个元素4135047与A中某元素匹配 → True
• ...

由于tf.math.equal无法直接比较不同形状的张量，以下是两种一次性实现的方法：

方法一：利用广播机制+维度归约 #

通过扩展张量形状实现广播比较，再对结果进行逻辑或归约：

import tensorflow as tf

# 定义张量（用tf.constant替代示例中的tf.Tensor，实际使用中可直接用已有张量）
A = tf.constant([4135047.0, 1193752.0], dtype=tf.float32)
B = tf.constant([1226019.0, 4135047.0, 4135047.0, 4169911.0, 1193752.0, 4135047.0, 4135047.0, 4135047.0], dtype=tf.float32)

# 扩展B的维度为(8,1)，与A广播后逐元素比较，得到(8,2)的布尔矩阵
matches = tf.math.equal(tf.expand_dims(B, axis=1), A)
# 对每个B元素对应的行取逻辑或，判断是否存在匹配
result = tf.math.reduce_any(matches, axis=1)

print(result)
# 输出：tf.Tensor([False  True  True False  True  True  True  True], shape=(8,), dtype=bool)

原理说明： #

1. tf.expand_dims(B, axis=1)将B的形状从(8,)转换为(8,1)，使得它能和形状为(2,)的A进行广播运算；
2. tf.math.equal会自动广播，生成一个(8,2)的布尔张量，每行对应B的一个元素与A所有元素的比较结果；
3. tf.math.reduce_any(matches, axis=1)对每行的布尔值取逻辑或，得到每个B元素是否在A中存在匹配的结果。

方法二：使用tf.math.in1d（简洁版） #

TensorFlow提供了直接实现该逻辑的APItf.math.in1d，无需手动处理形状：

import tensorflow as tf

A = tf.constant([4135047.0, 1193752.0], dtype=tf.float32)
B = tf.constant([1226019.0, 4135047.0, 4135047.0, 4169911.0, 1193752.0, 4135047.0, 4135047.0, 4135047.0], dtype=tf.float32)

result = tf.math.in1d(B, A)
print(result)
# 输出：tf.Tensor([False  True  True False  True  True  True  True], shape=(8,), dtype=bool)

原理说明： #

tf.math.in1d内部已处理了形状匹配问题，直接返回一个布尔张量，其中每个元素对应B中元素是否存在于A中，完全符合需求。

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