如何在GPU上无循环对二维张量每行应用torch.isin()？

现有两个二维张量：

tensor_1 = tensor([[10,  8,  7],
                   [ 4,  5,  9]])

target_tensor = tensor([[ 8,  8,  6, 11,  0],
                        [ 7,  2, 10,  4,  5]])

直接调用torch.isin(tensor_1, target_tensor)会得到全局匹配结果：

tensor([[ True,  True,  True],
        [ True,  True, False]])

但需要按每行分别执行isin匹配，得到行维度的精准结果：

tensor([[False,  True, False],
        [ True,  True, False]])

目前可通过for循环实现需求，但希望避免循环（尤其在GPU环境下提升计算效率）：

tmp_result = torch.zeros_like(tensor_1, dtype=torch.bool)

for i in range(tensor_1.size(0)):
    tmp_result[i] = torch.isin(tensor_1[i], target_tensor[i])

请问如何在GPU上不使用循环，对二维张量的每行应用torch.isin()？

可以通过张量维度扩展+广播比较+行维度逻辑聚合的方式实现，完全兼容GPU且无显式循环：

# 扩展tensor_1维度为 [批次大小, 元素数量, 1]，target_tensor为 [批次大小, 1, 每行目标数量]
expanded_tensor1 = tensor_1.unsqueeze(2)
expanded_target = target_tensor.unsqueeze(1)

# 广播后逐元素比较，再在目标元素维度取逻辑或，得到每行的匹配结果
result = (expanded_tensor1 == expanded_target).any(dim=2)

验证结果：

print(result)
# 输出：
# tensor([[False,  True, False],
#         [ True,  True, False]])

原理说明 #

1. unsqueeze扩展维度后，expanded_tensor1形状为(2,3,1)，expanded_target形状为(2,1,5)
2. 借助PyTorch广播机制，两个张量自动扩展为(2,3,5)的形状，逐元素比较每个元素是否存在于对应行的目标张量中
3. 最后在第2个维度（目标元素所在维度）上取any，只要存在匹配就返回True，最终得到行维度的匹配结果

这种方法完全基于张量并行操作，GPU环境下可充分利用硬件并行计算能力，效率远高于循环实现。

内容的提问来源于stack exchange，提问作者강태천