import numpy as np
import opt_einsum
print("opt_einsum version: ", opt_einsum.__version__)
no = 300
nv = 1200
naux = 50
nx = 2000
M = np.empty((nx,nx))
eta = np.empty((no,nx))
eta1 = np.empty((naux,nx))
theta = np.empty((naux,no))
subscripts = "iP,sP,PQ,jQ,tQ,ti,sj->"
print("\nOptimal path:")
path_info = opt_einsum.contract_path(
subscripts,
eta,eta1,M,eta,eta1,theta,theta,
optimize='optimal'
)
print(path_info[1])
print("\nGreedy path:")
path_info = opt_einsum.contract_path(
subscripts,
eta,eta1,M,eta,eta1,theta,theta,
optimize='greedy'
)
print(path_info[1])
opt_einsum version: 3.4.0
Optimal path:
Complete contraction: iP,sP,PQ,jQ,tQ,ti,sj->
Naive scaling: 6
Optimized scaling: 4
Naive FLOP count: 6.300e+15
Optimized FLOP count: 1.827e+11
Theoretical speedup: 3.449e+4
Largest intermediate: 6.000e+10 elements
--------------------------------------------------------------------------------
scaling BLAS current remaining
--------------------------------------------------------------------------------
3 0 PQ,sP->PQs iP,jQ,tQ,ti,sj,PQs->
4 0 PQs,jQ->PQsj iP,tQ,ti,sj,PQsj->
4 GEMM PQsj,sj->PQ iP,tQ,ti,PQ->
3 GEMM PQ,iP->Qi tQ,ti,Qi->
3 GEMM Qi,ti->Qt tQ,Qt->
2 DOT/EINSUM Qt,tQ-> ->
Greedy path:
Complete contraction: iP,sP,PQ,jQ,tQ,ti,sj->
Naive scaling: 6
Optimized scaling: 3
Naive FLOP count: 6.300e+15
Optimized FLOP count: 9.242e+8
Theoretical speedup: 6.817e+6
Largest intermediate: 4.000e+6 elements
--------------------------------------------------------------------------------
scaling BLAS current remaining
--------------------------------------------------------------------------------
3 GEMM ti,iP->tP sP,PQ,jQ,tQ,sj,tP->
3 GEMM sj,jQ->sQ sP,PQ,tQ,tP,sQ->
3 GEMM tP,tQ->PQ sP,PQ,sQ,PQ->
2 0 PQ,PQ->PQ sP,sQ,PQ->
3 GEMM PQ,sP->Qs sQ,Qs->
2 DOT/EINSUM Qs,sQ-> ->
Script to reproduce the problem:
Output: