Einstein summation in numpy

It turns out that numpy actually has several more mini-languages embedded in it. This next example is borrowed and slightly generalized from mathematics, where it is called Einstein summation.

Recall that matrix-matrix multiplication can be expressed by: $$ (AB)_{ij} = \sum_k A_{ik} B_{kj}$$

Einstein summation is a relatively natural way of generalizing this to arrays with multiple dimensions. The above matrix-matrix multiplication expression, for example, becomes:

$$ A_{ij} = B_{ik} C_{kj}$$

Where the implied rule is that repeated indices that are not part of the output will be summed over.

numpy simply takes this convention and turns it into a way of expressing array operations:

In [3]:
import numpy as np
In [4]:
A = np.random.randn(15, 20)

B = np.random.randn(20, 25)



AB1 = A.dot(B)
In [6]:
AB2 = np.einsum("ik,kj->ij", A, B)

print(np.linalg.norm(AB1 - AB2))
3.22729439922e-15