NumPy中np.einsum作用于张量的具体运算规则是什么？

np.einsum（爱因斯坦求和）是NumPy中实现张量缩并、维度重排、逐元素运算的通用函数，核心规则非常明确：

• 下标字符串按输入参数顺序，用单个字符对应每个输入张量的维度，字符和维度按位置一一匹配
• 相同字符标记的维度会被对齐做逐元素乘积，再沿该维度求和（即张量缩并操作）；同字符标记的维度长度必须相等，或其中一方长度为1触发广播
• ->右侧写下的字符，就是最终输出保留的维度，维度顺序和右侧字符的书写顺序完全一致
• 没有出现在->右侧的字符，不管是单个输入独有还是多个输入共有，最终都会沿该维度做求和操作，不会保留在输出结果中

第一组代码运算逻辑拆解 #

测试代码：

a = np.random.rand(2, 2, 5, 5)
b = np.random.rand(4, 5, 5, 1)
c = np.einsum('aijb,qwei->qweaj', b, a)

逐位对应维度与下标：

• 第一个输入为b，形状(4,5,5,1)，对应下标aijb：a维长4、i维长5、j维长5、b维长1
• 第二个输入为a，形状(2,2,5,5)，对应下标qwei：q维长2、w维长2、e维长5、i维长5

运算步骤：

1. 公共下标为i，两个输入的i维长度均为5，对齐后逐元素相乘，沿i维求和完成缩并
2. 下标b为输入b独有、未出现在输出侧，该维度长度为1，沿该维度求和后直接消去（长度为1的维度求和不会改变张量数值）
3. 剩余待保留维度按输出侧qweaj的顺序重排：q(2)、w(2)、e(5)、a(4)、j(5)，最终输出形状为(2,2,5,4,5)，和实际运行结果一致。

第二组代码运算逻辑拆解 #

测试代码：

a = np.random.rand(2, 2, 5, 4, 5)
b = np.random.rand(1, 2, 2, 4)
c = np.einsum('fqnd,qlkjd->nlkj', b, a)

逐位对应维度与下标：

• 第一个输入为b，形状(1,2,2,4)，对应下标fqnd：f维长1、q维长2、n维长2、d维长4
• 第二个输入为a，形状(2,2,5,4,5)，对应下标qlkjd：q维长2、l维长5、k维长2、j维长5、d维长4

运算步骤：

1. 公共下标为q和d，两个输入的q维长度均为2、d维长度均为4，对齐后逐元素相乘，沿q、d两个维度求和完成缩并
2. 下标f为输入b独有、未出现在输出侧，该维度长度为1，沿该维度求和后直接消去
3. 剩余待保留维度按输出侧nlkj的顺序重排：n(2)、l(5)、k(2)、j(5)，最终输出形状为(2,5,2,5)，和实际运行结果一致。

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