python - 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

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

Answer 1

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

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