11
$\begingroup$

Python's numpy has einsum function, which allows to express wide range of combinations of ND-arrays multiplications, transposings, convolutions etc in one short expression. It is based on Einstein's notation for tensor algebra.

It is much easier to unserstand einsum expression even for just matrix product.

In Mathematica I can't find any equivalent, and even simple linear matrix transform appeared in my code as follows:

modif = Sum[e[[chart, j]] * a[[j]], {j, 1, Dimensions[e][[2]]}];

A lot of indexings and complex Sum expression.

Is there still einsum equivalent in Mathematica? If there isn't, how could I implement a similarly convenient notation myself, without sacrificing too much performance?

UPDATE

The benefits of einsum it has in my opinion are following:

1) You should not think about order of operands and should not transpose them to correct shape.

For example, suppose you want to multiply matrix $A$ of (5,5) by vector of $X$ which you would write in math as $AX$. But suppose your vector $X$ is transposed to row (1, 5) which made $AX$ illegal.

You can compute

$(X(A^T))^T$

which can be easily miskaten as

$(XA)^T$

if you forget one transpose and it will work.

With einsum you would write

$a_{ij} x_{kj}$

where repeating of $j$ automatically means contraction and since it is in second position, it goes along row.

With einsum notaion you can even transpose result as you want:

np.einsum("ij,kj->ik", a, x)

will give column, while

np.einsum("ij,kj->ki", a, x)

will give row.

But you are not need to worry about it, because einsum can take any transpose.

2) You can easily deal with stacks of matrices and vectors and you should not think about how to extract them.

Suppose that now a is 10 matrices of compound shape of (10,5,5) and you want to multiple all of them by the same x and get 10 results stacked. Then I'd write

np.einsum("ijk,lk->ijl", a, x)

Or suppose x is also a stack of vectors and I want to multiply each matrix to corresponding x. The I would write

np.einsum("ijk,ilk->ijl", a, x)

Now the presense of i in output pattern also means no contraction.

3) einsum can be used to multiply more that 2 tensors in one expression.

You will NEVER use any other matrix operations, once you learn einsum and have efficient implementation of it!

$\endgroup$
  • 2
    $\begingroup$ I'm not sure if there is a single built-in that does what you want, but you can certainly make things easier by using Listability of Times to do Total[e[[chart, All]] a], which gives the same result as your sum (assuming a has the appropriate dimensions). $\endgroup$ – Kiro Oct 19 '17 at 11:23
  • 2
    $\begingroup$ Please include a number of example in the question that illustrate what einsum does. $\endgroup$ – Szabolcs Oct 19 '17 at 12:33
  • 2
    $\begingroup$ Have also a look at TensorProduct and TensorContract`. $\endgroup$ – Henrik Schumacher Oct 19 '17 at 12:37
  • 2
    $\begingroup$ I think some of your recent questions sound a bit link complaining that Mathematica doesn't have certain unusual numpy features (whether or not you meant them that way). But I also think that these questions could be made interesting and useful to many people. einsum is something so specific that it is very unlikely that any other system would have it. The answer to your literal question is that no, Mathematica doesn't have einsum. But writing such an answer is not very useful to anyone. That is why I edited your question and asked how to implement it. $\endgroup$ – Szabolcs Oct 19 '17 at 12:48
  • 2
    $\begingroup$ As you said, einsum is a "huge" function that can do many things. Imitating it exactly is not very exciting. But creating something comparable is. To get people's attention, and to get them excited, I suggest you describe a few basic but interesting use cases in the question itself, or even maybe show how you imagine a similar Mathematica function could be used (usage example, not implementation). It is important to include this in the question because most people won't click through to your link, and won't even realize that the question is interesting. Presentation really does matter :-) $\endgroup$ – Szabolcs Oct 19 '17 at 12:50
11
$\begingroup$

You can implement most of your einsum functionality using TensorContract/TensorTranspose. Here is an implementation, but note that it will not work with indices that are repeated but not contracted, and index specifications that don't match the corresponding array's depth:

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";

Here is your first example:

SeedRandom[1]
a = RandomReal[1, {3, 3}];
x = RandomReal[1, {3, 3}];

einsum[{{1,2}, {3,2}} -> {1,3}, a, x]

{{1.18725, 1.14471, 1.23396}, {0.231893, 0.203386, 0.416294}, {0.725267, 0.673465, 0.890237}}

Compare this with:

a . Transpose[x]

{{1.18725, 1.14471, 1.23396}, {0.231893, 0.203386, 0.416294}, {0.725267, 0.673465, 0.890237}}

Here is your second example:

einsum[{{1,2}, {3,2}} -> {3,1}, a, x]
x . Transpose[a]

{{1.18725, 0.231893, 0.725267}, {1.14471, 0.203386, 0.673465}, {1.23396, 0.416294, 0.890237}}

{{1.18725, 0.231893, 0.725267}, {1.14471, 0.203386, 0.673465}, {1.23396, 0.416294, 0.890237}}

Your third example:

SeedRandom[1];
a = RandomReal[1, {3,2,2}];
x = RandomReal[1, {2,2}];

einsum[{{1,2,3}, {4,3}} -> {1,2,4}, a, x]
a . Transpose[x]

{{{0.373209, 0.890669}, {0.380332, 0.926471}}, {{0.11833, 0.290096}, {0.286499, 0.720608}}, {{0.340815, 0.964971}, {0.274859, 0.824754}}}

{{{0.373209, 0.890669}, {0.380332, 0.926471}}, {{0.11833, 0.290096}, {0.286499, 0.720608}}, {{0.340815, 0.964971}, {0.274859, 0.824754}}}

As I said earlier, repeated indices that aren't contracted are not supported by my implementation, so your 4th example won't work. Finally, if you give einsum symbolic arrays, it will still work:

Clear[a, x]
$Assumptions = a ∈ Arrays[{10,5,5}] && x ∈ Arrays[{5,5}];

einsum[{{1,2,3}, {4,3}} -> {1,2,4}, a, x]

TensorContract[a \[TensorProduct] x, {{3, 5}}]

$\endgroup$
  • $\begingroup$ Hm, I need "ij,kl->ikjl" pattern and it doesn't seem to work: x = {{1, 1}, {1, 1}}; einsum[{{1, 2}, {3, 4}} -> {1, 3, 2, 4}, x, x] (* TensorTranspose::ttrank: Permutation {1,3,2,4} moves slots beyond tensor rank 2.*) $\endgroup$ – Yaroslav Bulatov Sep 26 at 16:18
  • 1
    $\begingroup$ @YaroslavBulatov It works for me in M12. Perhaps there's a version dependency? $\endgroup$ – Carl Woll Sep 26 at 16:26
  • $\begingroup$ Aha, seems to work fine in Mathematica Online, I may have redefined one too many system symbols $\endgroup$ – Yaroslav Bulatov Sep 26 at 16:41

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.