如何使用 numpy 将 N 个向量乘以 N 个矩阵？

Question

我有一个形状为 (N, L) 的矩阵 M 和一个形状为 (N, L, K) 的 3D 张量 P。我想获得形状为 (N, K) 的矩阵 V，其中 V[i] = M[i] @ P[i]。我可以用 for 循环来做，但效率很低，我想用一个或几个操作来做，这样它就可以在 CUDA 上并行 运行。

我试过像这样乘法

V = M @ P

但这会导致 3D 张量，其中 V[i, j] = M[j] @ P[i]。

np.diagonal(M @ P).T基本上是我想要的，但是这样计算很浪费计算量。

Answer 1

您可以使用 np.einsum:

>>> M = np.random.rand(5, 2)
>>> P = np.random.rand(5, 2, 3)

>>> V = np.einsum('nl,nlk->nk', M, P)

>>> V.shape
(5, 3)