如何使用PyTorch的tensordot()实现特定的点积向量计算？

你需要实现的是向量A的元素和与向量B各元素的乘积，最终得到与B同维度的结果向量。以下是用tensordot()实现的具体方案：

代码实现 #

首先定义输入张量，再通过指定tensordot()的轴匹配规则完成计算：

import torch

# 定义输入向量
A = torch.tensor([a1, a2])  # shape: (2,)
B = torch.tensor([b1, b2, b3])  # shape: (3,)

# 使用tensordot完成运算
result = torch.tensordot(A, B, dims=([], []))

结果验证 #

运行后result的shape为(3,)，元素完全符合你的预期：

[
  a1*b1 + a2*b1,
  a1*b2 + a2*b2,
  a1*b3 + a2*b3
]

参数说明 #

tensordot()的dims参数用于指定两个张量需要收缩的轴：

• 这里dims=([], [])表示不对任何轴做收缩，本质是计算A和B的外积后，对A的所有维度求和，最终得到与B同形状的结果。
• 如果想更直观地理解，也可以先给A增加一个维度再指定轴匹配，效果完全一致：

  result = torch.tensordot(A.unsqueeze(1), B, dims=([0], []))
  

等价替代方案 #

如果不限制必须用tensordot()，更简洁的写法是直接求和加元素乘法：

result = A.sum() * B

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