einsum

import numpy as np

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
 
A = np.array([
    [1, 2],
    [3, 4],
])
 
B = np.array([
    [5, 6],
    [7, 8],
])

Matrix multiplication

print(a @ b)
print(A @ B)

32
[[19 22]
 [43 50]]

print(np.einsum("i,i", a, b))
print(np.einsum("ij,jk->ik", A, B))

32
[[19 22]
 [43 50]]

Trace (Sum of product of diagnal values)

print(np.trace(A))
print(np.einsum("ii->", A))

5
5

Get diagnal

print(np.diag(A))
print(np.einsum("ii->i", A))

[1 4]
[1 4]

Sum over an axis

A

array([[1, 2],
       [3, 4]])

print(np.sum(A, axis=0))
print(np.einsum("ij->j", A))
 
## Under the hood
result = np.zeros((A.shape[1]), dtype=A.dtype)
for j in range(A.shape[1]):
    for i in range(A.shape[0]):
        result[j] = result[j] + A[i, j]
print(result)

[4 6]
[4 6]
[4 6]

View as result

np.einsum returns view of the array which is writable.

np.einsum("ii->i", A)[...] = 1
A

array([[1, 2],
       [3, 1]])

print(np.einsum("cd,da", A, B))
print(np.einsum("ij,jk", A, B))
 

[[19 43]
 [22 50]]
[[19 22]
 [43 50]]

Tensor contraction

$${\text{out}}_{kl}=\sum _{i,j}{a}_{ijk}\cdot b_{jil}$$

a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)
 
print(np.einsum('ijk,jil', a, b))
## Using loops
 
result = np.zeros((5, 2))
 
for k in range(a.shape[-1]):
    for l in range(b.shape[-1]):
        for i in range(a.shape[0]):
            for j in range(a.shape[1]):
                result[k, l] += a[i,j,k] * b[j, i, l]
print(result)

[[4400. 4730.]
 [4532. 4874.]
 [4664. 5018.]
 [4796. 5162.]
 [4928. 5306.]]
[[4400. 4730.]
 [4532. 4874.]
 [4664. 5018.]
 [4796. 5162.]
 [4928. 5306.]]

Optimizations

a = np.ones(64).reshape(2,4,8)

path = np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize="optimal")[0]
 
for iteration in range(10000):
    _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path)