4

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?

DathosPachy
  • 742
  • 1
  • 6
  • 17

1 Answers1

4

With np.einsum -

k = np.einsum('li,ijk,mi->lmjk',e,w,e)
Divakar
  • 218,885
  • 19
  • 262
  • 358