python - 仅沿最后两个维度获取矩阵乘积的单个操作

Question

经过反射(reflection)，这可能是显而易见的，但现在我还不清楚。

对于一对形状为 (K, N, M) 和 (K, M, N) 的 numpy 数组，用 a 表示和 b 分别，有没有办法将以下内容计算为单个向量化操作:

import numpy as np
K = 5
N = 2
M = 3
a = np.random.randn(K, N, M)
b = np.random.randn(K, M, N)

output = np.empty((K, N, N))
for each_a, each_b, each_out in zip(a, b, output):
    each_out[:] = each_a.dot(each_b)

一个简单的a.dot(b)返回每对第一个轴的点积(因此它返回一个形状的数组(K, N, K, N).

编辑:为那些无法理解问题的人稍微充实一下代码。

Answer 1

我不久前回答过类似的问题:Element-wise matrix multiplication in NumPy .

我认为您正在寻找的是:

output = np.einsum('ijk,ikl->ijl', a, b)

祝你好运!