CODEWITHSUNDEEP
pythonnumpy
Ulderique Demoitre
Ulderique Demoitre

Reputation: 1068

Speed up multilple matrix products with numpy

In python I have 2 three dimensional arrays:

T with size (n,n,n)

U with size (k,n,n)

T and U can be seen as many 2-D arrays one next to the other. I need to multiply all those matrices, ie I have to perform the following operation:

for i in range(n):
        H[:,:,i] =  U[:,:,i].dot(T[:,:,i]).dot(U[:,:,i].T)

As n might be very big I am wondering if this operation could be in some way speed up with numpy.

Upvotes: 2

Views: 77

Answers (1)

Divakar
Divakar

Reputation: 221504

Carefully looking into the iterators and how they are involved in those dot product reductions, we could translate all of those into one np.einsum implementation like so -

H = np.einsum('ijk,jlk,mlk->imk',U,T,U)

Upvotes: 3

Related Questions