Einsum Notation

linear-algebra

Definition

Einsum Notation

Einstein summation notation is a concise way to represent summations over repeated indices when performing tensor operations. Instead of writing out explicit summation symbols $\sum$, we reuse an index letter in multiple arrays to imply a summation over that index. This allows for compact expressions of operations like matrix multiplication, tensor contraction, outer product, etc.

Understanding

Let $A\in F^{m\times n}$ and $B\in F^{n\times 1}$ be two matrices:

$$\begin{array}{rl}A& =\left[\begin{array}{cccc}{a}_{11}& a_{12}& \dots & a_{1n}\\ a_{21}& a_{22}& \dots & a_{2n}\\ ⋮& ⋮& \ddots & ⋮\\ a_{m1}& a_{m2}& \dots & a_{mn}\end{array}\right]\\ B& =\left[\begin{array}{c}{b}_1\\ b_2\\ ⋮\\ b_n\end{array}\right]\end{array}$$

To compute $C=\mathrm{einsum}(‘ij,jk\to i{k}^{\prime },A,B)$, we need to replace the repeated indices $C_{ik}$. Since $j$ occurs on the left-hand side of the arrow and not on the right-hand side, we need to sum over $1\le j\le n$ for each $(i,k)$ pair.

The shape of the resulting tensor is equal to the dimensionalities of the indices $i$ and $k$, thus $m$ and $1$ respectively:

$$C=\left[\begin{array}{c}{c}_{11}\\ c_{21}\\ ⋮\\ c_{m1}\end{array}\right]\in F^{m\times 1}$$

Now, we need to compute each $c_{ik}$ element by summing over the $j$ index:

$$c_{ik}=\sum _{j=1}^n{a}_{ij}{b}_j$$

Thus, the resulting matrix $C$ is:

$$C=\mathrm{einsum}(‘ij,jk\to i{k}^{\prime },A,B)=\left[\begin{array}{c}{a}_{11}{b}_1+a_{12}{b}_2+\cdots +a_{1n}{b}_n\\ a_{21}{b}_1+a_{22}{b}_2+\cdots +a_{2n}{b}_n\\ ⋮\\ a_{m1}{b}_1+a_{m2}{b}_2+\cdots +a_{mn}{b}_n\end{array}\right]$$

Note that this is equivalent to matrix multiplication, which is a special case of tensor contraction.

Examples

The examples are taken from ¹.

Transpose

The transpose $B\in F^{n\times m}$ of a matrix $A\in F^{m\times n}$ is denoted as:

$$B_{ij}=A_{ji}$$

or equivalently:

einsum('ij->ji', A)

Sum

The sum of all elements of a matrix $A\in F^{m\times n}$ is denoted as:

$$A_{ij}=\sum _{i=1}^m\sum _{j=1}^n{A}_{ij}$$

or equivalently:

einsum('ij->', A)

Column Sum

The sum of all elements of the $j$-th column of a matrix $A\in F^{m\times n}$ is denoted as:

$$\sum _{i=1}^m{A}_{ij}$$

or equivalently:

einsum('ij->j', A)

Row Sum

The sum of all elements of the $i$-th row of a matrix $A\in F^{m\times n}$ is denoted as:

$$\sum _{j=1}^n{A}_{ij}$$

or equivalently:

einsum('ij->i', A)

Dot Product

The dot product of two vectors $x,y\in F^n$ is denoted as:

$$a_i{b}_i=\sum _{i=1}^n{a}_i{b}_i$$

einsum('i,i->', x, y)

Footnotes

1. https://rockt.ai/2018/04/30/einsum ↩