爱因斯坦求和约定einsum简单例题解读

概论

在爱因斯坦求和约定或 einsum() 格式字符串中，所有的索引都可以分为两类：自由索引集和求和索引集。它们的区别很简单：

• 自由索引是用于输出规范中的索引。它们与外层 for 循环相关联。
• 求和索引是所有其他索引：它们出现在参数规范中，但不出现在输出规范中。之所以称为求和索引，是因为在计算输出张量时，这些索引的维度被求和掉了。它们与内层 for 循环相关联。

根据这种划分，每类索引的角色就很明确了：自由索引用于遍历每个输出元素，而求和索引用于计算和求和所需的乘积项。

具体例子

import numpy as np
# 定义一个 2 × 2 矩阵
a = np.array([[2, 3],[4, 5]])

# ii 代表对角线向量(2, 5)
# ii,ij 代表 (2, 5) * ((2,3),(4,5))
# ij 是求和轴，ij都保留
np.einsum('ii,ij->ij', a, a)
输出：array([[ 4,  6],
       [20, 25]])

# 求和轴只有i，说明是沿着j轴求和
# 4+6=10,20+25=45，即（10,45)
np.einsum('ii,ij->i', a, a)
输出：array([10, 45])

#求和轴只有j，说明沿着i轴求和
#4+20, 6+25
np.einsum('ii,ij->j', a, a)
输出：array([24, 31])

#箭头右侧ij表示全部保留，只能原位计算了
np.einsum('ij,ij->ij', a, a)
输出：array([[ 4,  9],
       [16, 25]])