关于np.einsum维度变换逻辑及无einsum替代实现方案的咨询

原代码逻辑拆解 #

首先看你给出的代码：

import numpy as np
a = np.random.randint(0, 10, (2,3,4))
b = np.random.randint(0, 10, (3, 6, 4))
c = np.einsum('bld,hid->blhd', a,b)

einsum运算规则说明 #

einsum的下标规则决定了运算逻辑：

• 输入a的形状(2,3,4)和下标bld一一对应：b=2、l=3、d=4
• 输入b的形状(3,6,4)和下标hid一一对应：h=3、i=6、d=4
• 输出下标为blhd，仅输入下标i没有出现在输出中，该维度为求和维度：
  运算时先将a、b在公共维度d上对齐逐元素相乘，其余维度自动广播对齐，最后对i维度求和，得到的输出形状就是(b=2, l=3, h=3, d=4)，和你观测到的结果一致。

不使用einsum的实现方法 #

可以直接用Numpy的广播机制+sum操作实现完全等价的运算，代码如下：

import numpy as np

a = np.random.randint(0, 10, (2,3,4))
b = np.random.randint(0, 10, (3,6,4))

# 给a新增h、i两个占位维度，形状变为 (2,3,1,1,4)
a_expand = a[:, :, np.newaxis, np.newaxis, :]
# 给b新增b、l两个占位维度，形状变为 (1,1,3,6,4)
b_expand = b[np.newaxis, np.newaxis, :, :, :]
# 逐元素广播相乘后，对i维度（索引为3）求和得到结果
c_without_einsum = (a_expand * b_expand).sum(axis=3)

# 验证和einsum结果完全一致
c_einsum = np.einsum('bld,hid->blhd', a,b)
print(np.allclose(c_without_einsum, c_einsum)) # 输出为True代表结果完全等价

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