Numpy Einstein Summation

einsum.md

Einsum

A = np.array([[1, 1, 1],
              [2, 2, 2],
              [5, 5, 5]])

B = np.array([[0, 1, 0],
              [1, 1, 0],
              [1, 1, 1]])

np.einsum('ij,jk->ijk', A, B) 
array([[[0, 1, 0],
        [1, 1, 0],
        [1, 1, 1]],

       [[0, 2, 0],
        [2, 2, 0],
        [2, 2, 2]],

       [[0, 5, 0],
        [5, 5, 0],
        [5, 5, 5]]])

np.einsum('ij,jk->ik', A, B) 
array([[ 2,  3,  1],
       [ 4,  6,  2],
       [10, 15,  5]])

np.einsum('ij,jk->ij', A, B) 
array([[ 1,  2,  3],
       [ 2,  4,  6],
       [ 5, 10, 15]])

np.einsum('ij,jk->jk', A, B) 
array([[0, 8, 0],
       [8, 8, 0],
       [8, 8, 8]])

References

• ajcr - A basic introduction to NumPy's einsum
• Alex Rogozhnikov - einops
• Olexa Bilaniuk - Einstein Summation in Numpy
• Tim Rocktäschel - Einsum is All you Need