Einsum

Einstein Summation in Numpy or Pytorch

Key Ideas

• Element-wise multiplication is done along identical indexes in the inputs, e.g., "a,a->a".
• Summation is done along missing indexes in the output, e.g., "a->". (No output means scalar value)
• Transpositions in the output are automatically handled.

Examples

• Vector inner product: "a,a->"
• Vector element-wise product: "a,a->a"
• Vector outer product: "a,b->ab"
• 1-D summation: "a->"
• 2-D summation: "ab->"
• 3-D summation: "abc->"
• Matrix transposition: "ab->ba"
• Matrix diagonal: "ii->i"
• Matrix trace: "ii->"
• Matrix inner product: "ab,ab->"
• Matrix multiplication: "ab,bc->ac"
• Batch matrix multiplication: "Nab,Nbc->Nac"
• Quadratic form / Mahalanobis Distance: "a,ab,b->"

Documentation

• Pytorch: https://pytorch.org/docs/stable/generated/torch.einsum.html
• Numpy: https://numpy.org/doc/stable/reference/generated/numpy.einsum.html

Credits: https://obilaniu6266h16.wordpress.com/2016/02/04/einstein-summation-in-numpy/amp/