当输入是许多相同的数组时,使 np.einsum 更快? (或任何其他更快的方法)

Making a np.einsum faster when inputs are many identical arrays? (Or any other faster method)

我有一段代码类型:

nnt = np.real(np.einsum('xa,xb,yc,yd,abcde->exy',evec,evec,evec,evec,quartic))

其中 evec 是(比如)L x L np.float32 数组,quartic 是 L x L x L x L x T np.complex64 数组。

我发现这个例程比较慢。

我认为由于所有 evec 都是相同的,所以可能有更快的方法吗?

提前致谢。

首先,您可以重复使用第一个计算:

evec2 = np.real(np.einsum('xa,xb->xab',evec,evec))
nnt = np.real(np.einsum('xab,ycd,abcde->exy',evec2,evec2,quartic))

如果您不关心内存而只需要性能:

evec2 = np.real(np.einsum('xa,xb->xab',evec,evec))
nnt = np.real(np.einsum('xab,ycd,abcde->exy',evec2,evec2,quartic,optimize=True))