There are currently no responses for this story.

Be the first to respond.

When dealing with linear or multilinear algebra in Python, summation loops and NumPy functions can get quite messy, hard to read, and even slow. This was the case for me until I discovered NumPy's einsum function a while ago and I’m surprised not everyone is talking about it.

I am going to show you how to make your code more readable, concise, and efficient using Einstein notation in NumPy, TensorFlow, or PyTorch.

Understanding Einstein Notation

The basis of Einstein notation is to get rid of the summation symbol Σ when that doesn’t cause ambiguity (when we can determine the bounds of the indices).

Example #1: Product of matrices

In the following formula, the shape of the matrix A is (m, n) and the shape of B is (n, p).

Since we know the bounds for i, j, and k from the shapes of the matrices. We can simplify the formula to:

Example #2: Dot product of two vectors

The dot product of two n-dimensional vectors is:

We can write this in Einstein notation as:

Example #3: Dot product of two matrices

We can define a dot product of two matrices using this formula:

In Einstein notation, this is simply:

Example #4: Tensors

We can work with more than 2 indices. A tensor (higher-order matrix).

For example, we can write something like this:

Or even like this:

You get the idea!

When to use Einstein notation?

This mostly comes to when you’re working with vectors, matrices, and/or tensors, and you have to: multiply, transpose, and/or sum them in a particular way.

Writing the results of combining these operations can be simpler in Einstein notation.

Using Python’s einsum

einsum is implemented in numpy , torch , and tensorflow . In all of these modules, it follows the syntax einsum(equation, operands) .

Where we replace ■ by indices. And after -> we put the output indices.

This is equivalent to:

if an input or output is a scalar (it has no indices), we can leave the index empty.

Here are the examples above.

Example #1: Matrix multiplication

einsum("ik,kj->ij", A, B)

Example #2: Vector dot product

einsum("i,i->",u, v)

Example #3: Matrix dot product

einsum("ij,ij->", A, B)

Example #4: Tensors

einsum("ijkl,klij->ij", A, B)

einsum("iqrj,klqmr->ijklm", A, B)

You can use this with almost any formula involving linear algebra and multilinear algebra.

Performance

So how does einsum perform compared to using loops or numpy functions?

I decided to run example #3 using three methods:

After running 1,000,000 tests and using timeit :

• Loops: 24.36s
• Built-in functions: 7.58s
• Einsum: 3.78s

einsum is clearly faster. Actually, twice as fast as numpy’s built-in functions and, well, 6 times faster than loops, in this case.

Why is einsum fast?

This comes down to the fact that numpy is written in C.

When using native Python loops, all the data manipulation happens in the Python interpreter.

When using built-in numpy functions, it happens in C, which offers numpy developers the ability to optimize their code. This is why numpy is faster.

But when using einsum , numpy handles the data once in C and returns the final result, while using multiple numpy functions spends more time returning multiple values.

einsum can prove to be a great one-liner in some situations. While it is not only one way to improve the readability and efficiency of your code, it must be a no-brainer to use it when possible.

There are other ways to optimize Python code though, like using caching, which I am going to cover in a future article.