6
$\begingroup$

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]
$\endgroup$
  • $\begingroup$ I think your notations do not generalize well for larger number of indices. It seems the input can be encoded as {{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$ – yarchik Oct 4 at 17:28
  • 1
    $\begingroup$ I think your example in Mathematica would be represented as 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 such einsum specs to the above Tr + Dot representation be useful? $\endgroup$ – Carl Woll Oct 4 at 17:35
  • 1
    $\begingroup$ When each building block possesses only 2 indices the problem reduces to the well known question of finding a path that traverses all points of a graph visiting each point only one time. See FindHamiltonianPath $\endgroup$ – yarchik Oct 4 at 17:36
  • $\begingroup$ @CarlWoll yes it would be useful. The motivating example was to compute Xij Mijkl Xkl where Mijkl is written as (Mij Mkl + Mik Mjl + Mil Mjk) (this is Isserlis theorem), and I think it can be done in terms of Tr and Dot operators on rank-2 tensors $\endgroup$ – Yaroslav Bulatov Oct 4 at 18:02
  • $\begingroup$ @yarchik I don't see the Hamiltonian Path connection. For 2-index case, the problem reduces to finding triangulation of the graph that minimizes the size of largest maximal clique. Agree that using 1,2,3 is more convenient than "i","j","k" $\endgroup$ – Yaroslav Bulatov Oct 4 at 18:09
5
$\begingroup$

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

$\endgroup$
  • $\begingroup$ Hm I'm wondering if it could be made to work with X, M1, M2, X $\endgroup$ – Yaroslav Bulatov Oct 4 at 20:48
  • $\begingroup$ @YaroslavBulatov Something like $Assumptions = (X | M1 | M2) ∈ Matrices[{d, d}]; FromTensor@einsum[{{1,3},{1,2},{3,4},{2,4}}->{},X,M1,M2,X]? $\endgroup$ – Carl Woll Oct 4 at 20:59
  • $\begingroup$ Oh nice! I was able to check most of my derivations. This result of this one could be simpler FromTensor@einsum[{{1, 3}, {2, 4}, {1, 3}} -> {2, 4}, M, M, X] , I expected to see tr(M'X)M. Also couldn't related quantity FromTensor@einsum[{{1, 3}, {2, 4}, {1, 3}, {1, 4}} -> {}, M, M, X, X] $\endgroup$ – Yaroslav Bulatov Oct 4 at 21:17
  • $\begingroup$ @YaroslavBulatov Right, FromTensor (currently) doesn't work with tensors that have disconnected pieces. I'm planning to add support for that soon. In the meantime, you could use FromTensor @ TensorReduce @ ... instead. $\endgroup$ – Carl Woll Oct 4 at 21:19
  • $\begingroup$ Aha, that works! I think this should be part of standard Mathematica functionality, it could've saved me quite a bit of time $\endgroup$ – Yaroslav Bulatov Oct 4 at 21:33

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.