单一操作，仅沿最后两个维度获取矩阵产品

Question

这可能在反思中很明显，但我现在还不清楚。

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

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)

祝你好运！