2D和3D阵列之间的点积

Question

假设我有两个数组V和Q，其中V为(i, j, j)且Q为(j, j)。我现在希望计算Q的每个“行”V的点积，并将结果保存为(i, j, j)大小的矩阵。只需迭代i之类的

，就可以使用for循环轻松完成

import numpy as np

v = np.random.normal(size=(100, 5, 5))
q = np.random.normal(size=(5, 5))
output = np.zeros_like(v)

for i in range(v.shape[0]):
    output[i] = q.dot(v[i])

然而，这对我的需求来说太慢了，我猜测有一种方法可以使用einsum或tensordot来对这个操作进行矢量化，但我还没有想到它出。有人可以指点我正确的方向吗？感谢

Answer 1

你当然可以使用np.tensordot，但之后需要交换轴，如此 -

out = np.tensordot(v,q,axes=(1,1)).swapaxes(1,2)

使用np.einsum，它更直接，就像这样 -

out = np.einsum('ijk,lj->ilk',v,q)