Einstein summation, or einsum, is a powerful and flexible operation used for tensor manipulation. It allows you to perform summation over multiple dimensions concisely, making it a favourite among data scientists and machine learning practitioners.

Why Choose Einsum?

• Flexibility: Works with a variety of tensor shapes and operations.
• Conciseness: Eliminates the need for multiple intermediate steps like transpositions or reshaping.
• Simplicity: Reduces complex operations into a single line of code.

Einsum is available across major libraries like NumPy and PyTorch. Let’s explore its versatility using PyTorch, starting with some common use cases like matrix multiplication, summation, permutation, and more.

Syntax of Einsum

Einsum uses a simple notation to define operations:

• Indices: Represent dimensions (e.g., i, j, k).
• Commas: Separate tensors being operated on.
• Arrow (>): Indicates the desired output shape.

For example, 'ik, kj -> ij' means:

• The first tensor has dimensions ik.
• The second tensor has dimensions kj.
• The output should have dimensions ij, summing over the common dimension k.

Examples of Einsum Operations

Matrix Multiplication