高效计算PyTorch张量所有元素对差值的方法问询

对于大规模张量，嵌套循环完全不现实——PyTorch的广播机制和内置索引工具能轻松搞定这个需求，核心思路是先生成所有元素对的差值矩阵，再提取你需要的上三角区域（对应i<j的元素对）。

步骤示例 #

假设输入张量是：

import torch
x = torch.tensor([1, 2, 3, 4])

1. 扩展维度实现广播
   把原张量分别扩展成列向量和行向量，做差时会自动广播生成所有元素对的差值矩阵：

# 列向量：shape (4,1)
x_col = x.unsqueeze(1)
# 行向量：shape (1,4)
x_row = x.unsqueeze(0)
# 差值矩阵：shape (4,4)，其中diff_matrix[i][j] = x[i] - x[j]
diff_matrix = x_col - x_row

2. 提取上三角区域（排除对角线）
   我们需要的是i < j的元素对，对应差值矩阵的上三角部分（跳过对角线，设置offset=1）：

# 获取上三角的索引（i<j）
indices = torch.triu_indices(len(x), len(x), offset=1)
# 提取对应差值
result = diff_matrix[indices[0], indices[1]]

验证结果 #

打印result会得到：

tensor([-1, -2, -3, -1, -2, -1])

正好对应你要的[1-2, 1-3, 1-4, 2-3, 2-4, 3-4]。

为什么高效？ #

所有操作都是PyTorch底层优化的向量化计算，完全避开了Python循环的性能损耗，哪怕是几万甚至几十万元素的张量，速度也会比循环快几个数量级，还自动支持GPU加速（如果你的张量在GPU上）。

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