Tensor contraction with Einstein summation convention in Python

Tensor contraction with Einstein summation convention is a powerful technique for performing complex multi-dimensional array operations. Python's numpy.einsum() method provides an elegant way to implement tensor contractions using subscript notation.

Understanding Einstein Summation

The Einstein summation convention allows you to represent complex tensor operations using subscript labels. When indices appear in both input tensors but not in the output, they are automatically summed over (contracted).

Syntax

numpy.einsum(subscripts, *operands)

Parameters:

• subscripts ? String specifying the subscripts for summation as comma-separated labels
• operands ? Input arrays for the operation

Example: Tensor Contraction

Let's perform tensor contraction on two 3D arrays using the subscript pattern 'ijk,jil->kl' ?

import numpy as np

# Create two 3D arrays
arr1 = np.arange(60.).reshape(3, 4, 5)
arr2 = np.arange(24.).reshape(4, 3, 2)

print("Array1 shape:", arr1.shape)
print("Array2 shape:", arr2.shape)

# Perform tensor contraction using einsum
# 'ijk,jil->kl' means: contract over indices i and j, keep k and l
result = np.einsum('ijk,jil->kl', arr1, arr2)

print("\nResult shape:", result.shape)
print("Result (Tensor contraction):")
print(result)

Array1 shape: (3, 4, 5)
Array2 shape: (4, 3, 2)

Result shape: (5, 2)
Result (Tensor contraction):
[[4400. 4730.]
 [4532. 4874.]
 [4664. 5018.]
 [4796. 5162.]
 [4928. 5306.]]

How It Works

In the subscript 'ijk,jil->kl':

• ijk represents the first array with dimensions (i=3, j=4, k=5)
• jil represents the second array with dimensions (j=4, i=3, l=2)
• kl specifies the output dimensions (k=5, l=2)
• Indices i and j are contracted (summed over)

Common Use Cases

Conclusion

NumPy's einsum() provides a concise way to perform tensor contractions using Einstein summation notation. The subscript pattern defines which dimensions to contract and which to keep in the output, making complex tensor operations more readable and efficient.