Alternatives to numpy.einsum

Question

Given an $n_1 \times \cdots \times n_k \times g \times g$ tensor $A$ (i.e. a collection of $g \times g$ matrices) and an $n_1 \times \cdots \times n_k \times g$ tensor $b$ (i.e. a collection of vectors of length $g$) I want to use Numpy to compute the $n_1 \times \cdots \times n_k \times g$ tensor $c$ where each element is the matrix-vector product

$$c_{i_1,\ldots,i_k} = A_{i_1,\ldots,i_k}b_{i_1,\ldots,i_k}.$$

I found a way to do this using numpy.einsum but was curious if there was another way to do this or if there is a particular name for this operation? If there is an alternative how does the performance of this alternative compare to using einsum?

Below is an example:

from numpy import dot, einsum, zeros_like
from numpy.linalg import norm
from numpy.random import randn

n = 10
g = 4

matrices = randn(n,n,g,g)
vectors = randn(n,n,g)

# "manual" mat-vec multiplication (slow)
out1 = zeros_like(vectors)
for i in range(n):
    for j in range(n):
        out1[i,j] = dot(matrices[i,j], vectors[i,j])

# using numpy.ensum (faster)
out2 = zeros_like(vectors)
out2 = einsum('...ij,...j->...i', matrices, vectors)

numpy.isclose(out1, out2).all()  # outputs `True`

Answer 1

Pyhton 3.5 will introduce a new operator @, which was proposed by NumPy devs to be the matrix multiplication operator. You may want to read PEP 465, but the syntax that NumPy will adopt for this operator is "stacked matrix multiplication," i.e. exactly what you want.

There may be bits and pieces missing, but some work is already in, see the relevant PR, and I think that if you build the latest NumPy development version against a beta release of Python 3.5, you should be able to simply do something like:

out = matrices @ vectors

or perhaps:

out = matrices @ vectors[..., None]

to add a size-1 dimension to the end of your vectors tensor.

With earlier versions of Python, if you build the latest development version of NumPy, you should be able to do:

import numpy as np

out = np.linalg.matmul(matrices, vectors)  # or vectors[..., None]

These should be very fast, as they combine using a BLAS library to perform the matrix multiplications if available, with C iteration over all the stacked matrices.

If you don't want to get into non stable release territory (although NumPy 1.10 should be out the door in a few weeks), you have to choose:

• einsum has C iteration speed, but the actual matrix multiplication cannot match BLAS, even though it uses SIMD to be faster than a naive C implementation. Probably best for big n and small g.
• iteration plus dot is slow iterating, but fast multiplying (if you have a good BLAS library), so best for small n and large g.