python – Numpy矩阵乘法广播

我有一个由N 3×3数组组成的数组(矩阵的集合,虽然数据类型是np.ndarray),我有一个由N 3×1数组(向量集合)组成的数组.我想要做的是将每个矩阵乘以每个向量,所以我希望得到N个3×1数组.

简单的例子:

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广播?这个操作最终需要完成数千次,而且我不希望地图或列表理解操作减慢速度.

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

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