向量化NumPy中的多个向量 - 矩阵乘法

Question

我无法确定如何安排我的坐标轴，以便我可以以矢量化的方式执行以下操作。

基本上我有向量的阵列，矩阵的阵列，并且我要评估VMV^T为每个对应的矢量V和矩阵M

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速度100X：

%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