Python's einsum equivalent in Mathematica?

Question

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!

Answer 1

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}}]

Answer 2

There is now a Resource Function that implements Einsum. https://resources.wolframcloud.com/FunctionRepository/resources/EinsteinSummation/