Use of `np.einsum()`

np.einsum() is a powerful function in NumPy that performs Einstein summation, which allows for flexible manipulation of multi-dimensional arrays (tensors) using summation notation. It can handle various operations like matrix multiplication, element-wise operations, and tensor contractions in a very efficient way.

Syntax:

np.einsum(subscripts, *operands, **kwargs)

• subscripts: A string representing the Einstein summation convention.
• *operands: The arrays (tensors) on which the operation is performed.
• **kwargs: Optional arguments like optimize, which can be used to improve performance.

Einstein Summation Convention

The Einstein summation convention is a notational shorthand where repeated indices are implicitly summed over. For example:

• ij,jk->ik denotes a matrix multiplication between two 2D arrays.
• ii->i sums the diagonal elements of a matrix.

Common Examples

1. Matrix Multiplication (np.dot or np.matmul)

   import numpy as np
   A = np.array([[1, 2], [3, 4]])
   B = np.array([[5, 6], [7, 8]])
   # Matrix multiplication
   result = np.einsum('ij,jk->ik', A, B)
   print(result)
   

   Explanation:

   • ij: Refers to the indices of the matrix A.
   • jk: Refers to the indices of the matrix B.
   • ik: The result is the matrix multiplication of A and B.

2. Sum over an axis (similar to np.sum)

   input_array = np.array([[1, 2], [3, 4]])
   # Sum over all elements (similar to np.sum)
   total_sum = np.einsum('ij->', input_array)
   print(total_sum)
   

   Explanation:

   • ij->: This notation sums over all indices of the matrix input_array, returning the total sum.

3. Trace of a Matrix (sum of diagonal elements)

   input_array = np.array([[1, 2], [3, 4]])
   # Trace (sum of diagonal elements)
   trace = np.einsum('ii->', D)
   print(trace)
   

   Explanation:

   • ii->: This notation picks the diagonal elements of the matrix input_array and sums them.

4. Element-wise multiplication

   first_array = np.array([[1, 2], [3, 4]])
   second_array = np.array([[5, 6], [7, 8]])
   # Element-wise multiplication
   element_wise = np.einsum('ij,ij->ij', first_array, second_array)
   print(element_wise)
   

   Explanation:

   • ij,ij->ij: The indices ij are the same for both arrays first_array and second_array, resulting in element-wise multiplication.
5. Dot product of vectors (np.dot)

   x = np.array([1, 2, 3])
   y = np.array([4, 5, 6])
   # Dot product
   dot_product = np.einsum('i,i->', x, y)
   print(dot_product)
   

6. Tensor contraction (generalized summation over axes)

   first_array = np.random.rand(3, 3, 3)
   second_array = np.random.rand(3, 3)
   # Contract tensor G and matrix H
   tensor_contraction = np.einsum('ijk,jl->ikl', first_array, second_array)
   print(tensor_contraction.shape)
   

   Explanation:

   • ijk,jl->ikl: Summation is performed over the common axis j, resulting in a contraction of the tensor.

     Advantages of np.einsum

• Flexibility: You can perform many types of operations in a single function call.
• Efficiency: It can be faster than separate functions like np.dot, np.sum, etc., especially when you have complex operations to perform.