Given an einsum like below, how could I generate an efficient computation graph for it?
$$X_{ik} M_{ij}M_{kl} X_{jl}$$
The indices range from $1$ to $d$ and the goal is to minimize computation time assuming $d$ is large. IE, prefer $O(d^{k})$ to $O(d^{k+1})$. For the sum above, it can be computed as follows:
$$A_{kj}=X_{ik} M_{ij}\\B_{kj} = M_{kl} X_{jl}\\c=A_{kj}B_{kj}$$
You could specify this solution in terms of indices occurring in the expression
A={ik,ij}
B={kl,jl}
c={A,B}
More compactly, the problem and solution can be encoded as follows:
input: {ik, ij, kl, jl}
output: {{ik, ij}, {kl, jl}}
This is likely to be an NP-complete problem, but there are probably heuristics to find near-optimal solution most of the time.
Edit: the most important case for practical applications was when result can be expressed in terms matrices, which can be done using Carl Woll's package in the answer. Specifically, it seems to work to get efficient matrix expression for the following einsum
$$X_{ik} (M_{ij}^{(1)} M_{kl}^{(2)} + M_{ik}^{(3)} M_{jl}^{(4)} + M_{il}^{(5)} M_{jk}^{(6)}) X_{jl}$$
as
$$\text{tr}(M_2' X' M_1 X)+\text{tr}(M_3' X)\text{tr}(M_4' X)+\text{tr}(M_6' X M_5' X)$$
This was computed using the answer below as
PacletInstall[
"TensorSimplify",
"Site" -> "http://raw.githubusercontent.com/carlwoll/TensorSimplify/master"
]
<< TensorSimplify`
einsum[in_List -> out_, arrays__] :=
Module[{res = isum[in -> out, {arrays}]}, res /; res =!= $Failed];
isum[in_List -> out_, arrays_List] :=
Catch@Module[{indices, contracted, uncontracted, contractions,
transpose},
If[Length[in] != Length[arrays],
Message[einsum::length, Length[in], Length[arrays]];
Throw[$Failed]];
MapThread[
If[IntegerQ@TensorRank[#1] && Length[#1] != TensorRank[#2],
Message[einsum::shape, #1, #2];
Throw[$Failed]] &, {in, arrays}];
indices = Tally[Flatten[in, 1]];
If[DeleteCases[indices, {_, 1 | 2}] =!= {},
Message[einsum::repeat,
Cases[indices, {x_, Except[1 | 2]} :> x]];
Throw[$Failed]];
uncontracted = Cases[indices, {x_, 1} :> x];
If[Sort[uncontracted] =!= Sort[out],
Message[einsum::output, uncontracted, out];
Throw[$Failed]];
contracted = Cases[indices, {x_, 2} :> x];
contractions = Flatten[Position[Flatten[in, 1], #]] & /@ contracted;
transpose = FindPermutation[uncontracted, out];
Activate@
TensorTranspose[
TensorContract[Inactive[TensorProduct] @@ arrays, contractions],
transpose]]
einsum::length =
"Number of index specifications (`1`) does not match the number of \
arrays (`2`)";
einsum::shape =
"Index specification `1` does not match the array depth of `2`";
einsum::repeat =
"Index specifications `1` are repeated more than twice";
einsum::output =
"The uncontracted indices don't match the desired output";
$Assumptions = (X | M | M1 | M2 | M3 | M4 | M5 | M6) \[Element]
Matrices[{d, d}];
FromTensor@einsum[{{1, 3}, {1, 2}, {3, 4}, {2, 4}} -> {}, X, M1, M2, X]
FromTensor@
TensorReduce@
einsum[{{1, 3}, {2, 4}, {1, 3}, {2, 4}} -> {}, M3, M4, X, X]
FromTensor@
TensorReduce@
einsum[{{1, 4}, {2, 3}, {1, 3}, {2, 4}} -> {}, M5, M6, X, X]
{{1,3},{1,2},{3,4},{2,4}}
and the output as{{{1,3},{1,2}},{{3,4},{2,4}}}
allowing further generalizations for an arbitrary number of indices. $\endgroup$Tr[Transpose[M].Transpose[X].M.X]
, and I think evaluating this as normal would have the desired complexity. If so, would having a way to convert sucheinsum
specs to the aboveTr
+Dot
representation be useful? $\endgroup$FindHamiltonianPath
$\endgroup$