在NumPy中矢量化多个向量矩阵乘法

Question

我无法弄清楚如何排列轴,以便我可以以矢量化方式执行以下操作.

基本上我有一个向量数组,一个矩阵数组,我想评估每个对应的向量V和矩阵M的VMV ^ T

import numpy as np

N = 5 # normally 100k or so
vecs = np.random.rand(N, 2)
mats = np.random.rand(N, 2, 2)

output = np.array([np.dot(np.dot(vecs[i, ...], mats[i, ...]), vecs[i, ...].T) for i in range(N)])

如果它更简单,矢量化下面的中间结果也会有所帮助:

intermediate_result = np.array([np.dot(vecs[i, ...], mats[i, ...]) for i in range(N)])
# then I can do
output = np.sum(intermediate_result * vecs, axis=-1)

Answer 1

einsum对于N = 100k,基于An 的解决方案比你的循环快100倍:

%timeit np.array([np.dot(np.dot(vecs[i, ...], mats[i, ...]), vecs[i, ...].T) for i in range(N)])
%timeit np.einsum('...i,...ij,...j->...', vecs, mats, vecs)
np.allclose(np.array([np.dot(np.dot(vecs[i, ...], mats[i, ...]), vecs[i, ...].T) for i in range(N)]),
            np.einsum('...i,...ij,...j->...', vecs, mats, vecs))
1 loops, best of 3: 640 ms per loop
100 loops, best of 3: 7.02 ms per loop
Out[45]: True