I am trying to calculate the following using numpy's einsum:
import numpy as np
tmp_ee = np.ones((2,4,4))
tmp_ij = np.ones((2,2,4,2,4,2))
print(tmp_ee.shape)
print(tmp_ij.shape)
np.einsum('naq,nnbpqp->nab', tmp_ee, tmp_ij, optimize=True)However, I run into the following error:
ValueError: dimensions in operand 0 for collapsing index 'q' don't match (4 != 2)
If you look at my code above, the index 'q' has dimensions equal to 4 in both arrays. The other repeating indices also have consistent dimensions. Hence, I do not understand the cause of this error. In addition to this, when my second array temp_ij has dimensions of the form (2,2,3,3,3,3) or (2,2,1,1,1,1) or (2,2,2,2,2,2), i.e., (2,2,x,x,x,x) then the error above does not come up and the code runs smoothly. I tried to work around this by doing the operation in two steps using einsum, as follows:
tmp_ee = np.ones((2,4,4))
tmp_ij = np.ones((2,2,4,2,4,2))
print(tmp_ee.shape)
print(tmp_ij.shape)
temp = np.einsum('naq,nnbrqp->nabrp', tmp_ee, tmp_ij, optimize=True)
np.einsum('nabpp->nab', temp, optimize = True)This does not give an error like the code before did. This works fine and gives me the result. While this workaround appears to be working fine, it has slowed down my code significantly because I have to do this computation millions of times in my code. Is there an explanation for the error above?
Thanks!