Einops in 1 Minute

Learn to use `einops` for concise, readable, beautiful machine learning code.

This article was originally published at fkodom.substack.com.

Photo by Collab Media on Unsplash

Machine learning involves a lot of matrix operations. einops makes them easier and (hopefully) more intuitive.

Note: einops is compatible with NumPy, PyTorch, TensorFlow, and JAX. I use PyTorch in this article.

Matrix Multiplication

The formula for matrix multiplication is:

So, two matrices with indices (i,k) and (k,j) combine to make a matrix with indices (i,j). In shorthand notation,

i k, k j -> i j

This is Einstein notation! Here’s how to do matrix multiplication with einops:

Reshape Arrays

To swap the positions of axes 1 and 2:

To add a new axis at position 1:

Attention Mechanism

As a challenge problem, let’s implement a basic attention mechanism! The formula for attention is:

Written in PyTorch with einops: