用 einsum 进行张量乘法

Question

我有一个张量 phi = np.random.rand(n, n, 3) 和一个矩阵 D = np.random.rand(3, 3)。我想沿着 phi 的最后一个轴乘以矩阵 D，以便输出的形状为 (n, n, 3)。我试过这个

np.einsum("klj,ij->kli", phi, D)

但是我对这个表示法一点信心都没有。基本上我想做

res = np.zeros_like(phi)           
for i in range(n):
    for j in range(n):
        res[i, j, :] = D.dot(phi[i, j, :])

Answer 1

您将 phi 视为一个 n, n 向量数组，每个向量左乘 D。所以你想保持形状的 n, n 部分完全原样。向量的最后一个（唯一）维度应该与矩阵的最后一个维度相乘并相加（向量隐式为 3x1）：

np.einsum('ijk,lk->ijl', phi, D)

或

np.einsum('ij,klj->kli', D, phi)

通过 np.matmul（@ 运算符）使用广播可能更简单：

np.squeeze(D @ phi[..., None])

如果您不介意末尾的额外单位维度，您可以省略 squeeze。