对角线阵列上的numpy线性代数，没有明显的重复

Question

我有一个数组w （形状(3000, 100, 100) ），我想与另一个数组e （形状(5, 3000) ）相乘，使得结果k具有形状(5, 5, 100, 100)和

k[:, :, i, j] = e @ np.diag(w[:, i, j]) @ e.T

由于w太大，因此制作一些形状为(3000, 3000, 100, 100) super_w数组并明确填充主对角线是不实际的。 循环遍历i和j也不是非常有效。 有没有一种内存有效的方法来做到这一点，除了将w分解成块？

Answer 1

使用np.einsum -

k = np.einsum('li,ijk,mi->lmjk',e,w,e)