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]
Answer
Maybe the following will be useful for you.
You can combine my FromTensor function (part of my TensorSimplify paclet) with my einsum function to convert your einsum representation into Tr + Dot.
$Assumptions = (X|M) ∈ Matrices[{d,d}];
FromTensor @ einsum[{{1,3}, {1,2}, {3,4}, {2,4}}->{}, X, M, M, X]
Tr[Transpose[M].Transpose[X].M.X]
Hopefully the loading instructions for these functions is clear from the above links. If not, I can add them here again.
Addendum
If your tensor has disconnected pieces, then FromTensor doesn't currently work. A simple fix is to include TensorReduce. From the comments in the examples (I think I fixed a typo in the second example):
$Assumptions = (X | M) ∈ Matrices[{d,d}];
FromTensor @ TensorReduce @ einsum[{{1, 3}, {2, 4}, {1, 3}} -> {2, 4}, M, M, X]
FromTensor @ TensorReduce @ einsum[{{1, 3}, {2, 4}, {1, 3}, {2, 4}} -> {}, M, M, X, X]
M Tr[Transpose[M].X]
Tr[Transpose[M].X]^2
Comments
Post a Comment