numpy linear algebra on diagonal arrays without explicit duplication

Question

I have an array w (shape (3000, 100, 100) ) which I want to multiply with another array e (shape (5, 3000) ) such that the result k has shape (5, 5, 100, 100) and

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

Since w is so large, it's not practical to make some super_w array with shape (3000, 3000, 100, 100) and explicitly populate the main diagonal. It is also not terribly efficient to loop over i and j . Is there a memory-efficient way to do this, other than breaking up w into chunks?

Answer 1

使用np.einsum -

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