I am trying to find the fastest way to contract multidimensional arrays.
Basically, a more efficient way to do this:
import numpy as np def naive_contraction(tensor, indices): result = 0.0 for i in range(5): for j in range(5): for k in range(5): result += tensor[i,j,k] * tensor[indices, indices, i] * tensor[indices, indices, j] * tensor[indices, indices, k] return result
Then we simply perform the contraction like this:
tensor = np.random.rand(5,5,5) indices = [0,1,4,3,0,1] result = naive_contraction(tensor, indices)
Notice that the index values must be received as input. We don’t want to store the full multidimensional array obtained with all possible combinations of indices as it would be way too large (this is just a naive example to explain the situation).
It’s my understanding that Numpy stores data in row-major order. Therefore I put the “summation index” along the last axis of the multidimensional array wherever possible. How can I vectorize (and/or optimize) a tensor contraction like this in Python?
Thank you so much,
all the best.