Numpy 矩阵乘法广播

Question

我有一个由 N 个 3x3 数组（矩阵集合，尽管数据类型是 np.ndarray）组成的数组，我有一个由 N 个 3x1 数组（向量集合）组成的数组。我想要做的是将每个矩阵乘以每个向量，所以我希望得到 N 个 3x1 数组。

简单的例子：

A = np.ones((6,3,3))
B = np.ones((6,3,1))
np.dot(A,B) # This gives me a 6x3x6x1 array, which is not what I want
np.array(map(np.dot,A,B)) # This gives me exactly what I want, but I don't want to have to rely on map

我已经厌倦了各种重塑、探索einsum等，但无法让它按照我想要的方式工作。我如何让它与 numpy 广播一起使用？这个操作最终需要执行数千次，我不希望map或列出理解操作来减慢速度。

Answer 1

您可以使用np.einsum计算点积并创建所需形状的矩阵：

np.einsum('ijk,ikl->ijl', A, B)