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

Question

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

我试着像这样乘它

V = M @ P

但这会产生一个 3D 张量，其中V[i, j] = M[j] @ P[i].

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

标签： pythonnumpyreshape

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)