Einops in 1 Minute

Jul 23, 2023

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:

\(A_{ij} = \sum_{k} B_{ik}C_{kj}\)

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:

To swap the positions of axes 1 and 2:

To add a new axis at position 1:

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

\(\text{Attention}(q, k, v) = \text{softmax} \left( \frac{q \, k^T}{\sqrt{d_k}} \right) v\)

Written in PyTorch with einops:

Thank you for reading Frank Odom | Substack. This post is public so feel free to share it.