torch.einsum

爱因斯坦简记法：是一种由爱因斯坦提出的，对向量、矩阵、张量的求和运算∑的求和简记法。
这样的求和简记法，能够以一种 统一的方式表示各种各样的张量运算（内积、外积、转置、点乘、矩阵的迹、其他自定义运算）， 为不同运算的实现提供了一个统一模型。

print(a_tensor)
 
tensor([[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]])
 
print(b_tensor)
 
tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3],
        [4, 4, 4, 4]])
 
# 'ik, kj -> ij'语义解释如下：
# 输入a_tensor: 2维数组，下标为ik,
# 输入b_tensor: 2维数组，下标为kj,
# 输出output：2维数组，下标为ij。
# 隐含语义：输入a,b下标中相同的k，是求和的下标，对应上面的例子2的公式
output = torch.einsum('ik, kj -> ij', a_tensor, b_tensor)
 
print(output)
 
tensor([[130, 130, 130, 130],
        [230, 230, 230, 230],
        [330, 330, 330, 330],
        [430, 430, 430, 430]])

上面的例子指a的各行分别与b的各列相乘求和。

参考

https://blog.csdn.net/a2806005024/article/details/96462827
https://www.cnblogs.com/mengnan/p/10319701.html