如何理解Einsum表达式abc,cde->abde？求等效矩阵乘法代码

核心逻辑拆解 #

abc,cde->abde 的本质是：

• 将第一个张量 abc（维度：(a,b,c)）视为 a*b 个长度为 c 的向量
• 将第二个张量 cde（维度：(c,d,e)）视为 d*e 个长度为 c 的向量
• 对每一对来自第一个张量的向量和第二个张量的向量做内积（对应元素相乘后求和，即矩阵乘法的核心操作），最终得到的(a*b, d*e)矩阵再重塑为(a,b,d,e)的四维张量。

等效代码实现（基于 NumPy） #

先定义测试用的随机张量：

import numpy as np

# 定义各维度的具体大小，可按需调整
a, b, c, d, e = 2, 3, 4, 5, 6
tensor_abc = np.random.rand(a, b, c)
tensor_cde = np.random.rand(c, d, e)

1. einsum 基准实现

result_einsum = np.einsum('abc,cde->abde', tensor_abc, tensor_cde)

2. 等效矩阵乘法分步实现

# 步骤1：将三维张量 abc 重塑为 (a*b, c) 的二维矩阵
mat_ab_c = tensor_abc.reshape(a*b, c)
# 步骤2：将三维张量 cde 重塑为 (c, d*e) 的二维矩阵
mat_c_de = tensor_cde.reshape(c, d*e)
# 步骤3：执行标准矩阵乘法，得到 (a*b, d*e) 的中间结果
mat_ab_de = mat_ab_c @ mat_c_de
# 步骤4：将二维矩阵重塑为目标四维张量
result_equiv = mat_ab_de.reshape(a, b, d, e)

3. 验证结果一致性

通过以下代码确认两种实现的结果完全一致：

print(np.allclose(result_einsum, result_equiv))  # 输出 True 表示结果匹配

直观循环实现（用于理解逻辑，效率较低） #

如果想更清晰地对应到逐元素操作，可以用嵌套循环实现（仅作理解用，实际场景优先用矩阵乘法或einsum）：

result_loop = np.zeros((a, b, d, e))
for ai in range(a):
    for bi in range(b):
        for di in range(d):
            for ei in range(e):
                # 对维度c做内积求和，对应einsum的核心操作
                result_loop[ai, bi, di, ei] = np.dot(tensor_abc[ai, bi, :], tensor_cde[:, di, ei])

验证循环实现与einsum结果一致：

print(np.allclose(result_einsum, result_loop))  # 输出 True

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