Summary

The article discusses the use of Einstein notation and the einsum function in NumPy, TensorFlow, or PyTorch to make Python code more readable, concise, and efficient when dealing with linear or multilinear algebra.

Abstract

The article introduces the concept of Einstein notation, which simplifies mathematical expressions by eliminating the summation symbol when the bounds of indices are known. It provides examples of using Einstein notation for matrix multiplication, vector dot product, matrix dot product, and tensors. The article then explains how to use Python's einsum function, which follows the syntax einsum(equation, operands), to achieve the same results as Einstein notation. The article concludes by comparing the performance of einsum to using loops or built-in numpy functions, showing that einsum is faster and more efficient.

Opinions

The author believes that using Einstein notation and einsum can make Python code more readable, concise, and efficient.
The author is surprised that not everyone is using einsum for linear or multilinear algebra in Python.
The author suggests that einsum can be a great one-liner in some situations.
The author recommends using einsum when possible to improve the readability and efficiency of code.
The author acknowledges that there are other ways to optimize Python code, such as using caching.
The author promotes the use of an AI service that provides the same performance and functions as ChatGPT Plus(GPT-4) but is more cost-effective.
The author provides a special offer for the AI service at $1/month.

Write Better And Faster Python Using Einstein Notation

Make your code more readable, concise, and efficient using “einsum”

Photo by Lewis Kang'ethe Ngugi on Unsplash

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.