Numpy linear algebra on diagonal arrays without explicit duplication

Question

I have an array w(shape (3000, 100, 100)) that I want to multiply by another array e(shape (5, 3000)), so the result khas a shape (5, 5, 100, 100)and

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

Since it is wso large, it is impractical to make some array super_wwith a form (3000, 3000, 100, 100)and explicitly fill the main diagonal. It is also not very efficient to iterate over iand j. Is there a memory efficient way to do this other than breaking winto pieces?

Answer 1

C np.einsum-

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