如何对两个3D张量沿第一维度执行矩阵乘法运算？

你遇到的报错是因为np.einsum的语法规则要求输出子串中的每个下标只能出现一次，重复下标仅允许出现在输入子串中，代表对该维度执行规约求和。你写的->ijj中j出现了两次，相当于试图给输出的两个不同维度分配同一个下标，不符合语法要求。

你需要计算的本质是沿第一维对齐的批量二维矩阵乘法：每个batch下，形状为(5,2)的矩阵和(2,5)的矩阵相乘得到(5,5)的结果，最终拼接为(3,5,5)的输出，以下是几种实现方式：

1. 修正后的einsum写法 #

只需要给b的最后一维分配独立的下标即可：

import numpy as np

a = np.random.randn(30).reshape(3, 5, 2)
b = np.random.randn(30).reshape(3, 2, 5)

# 输出形状为 (3,5,5)，符合预期
res = np.einsum('ijk,ikl->ijl', a, b)

逻辑说明：

• 下标i对应两个张量的第一维（batch维度），保留且不做求和
• 下标k对应a的第三维、b的第二维，两个维度长度均为2，执行点积求和
• 下标j对应a的第二维，下标l对应b的第三维，保留为输出的后两维

2. 更简洁的批量矩阵乘法实现 #

numpy的@运算符（底层调用np.matmul）原生支持带batch维度的矩阵乘法，写法更简单：

res = a @ b
# 等价于显式调用
res = np.matmul(a, b)

3. tensordot实现说明 #

如果一定要用np.tensordot实现，需要注意默认tensordot不会自动对齐未指定的公共维度，直接写np.tensordot(a, b, axes=([2], [1]))会得到(3,5,3,5)的冗余维度，需要额外做维度筛选，因此这类带batch对齐的乘法场景更推荐用einsum或者matmul实现。

• 如果需要对其他维度执行点积，只需要调整einsum的下标对应关系即可，灵活度非常高
• 如果确实需要对重复下标维度做规约（比如取矩阵的迹），可以省略输出子串，einsum会自动对所有重复下标求和，比如np.einsum('ijj->i', arr)可以得到每个batch矩阵的迹

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