Evaluate the lowest cost contraction order for an einsum expression in Python

To find the optimal contraction order for an einsum expression, use numpy.einsum_path(). This function analyzes different ways to perform the tensor contractions and returns the most efficient path to minimize computational cost.

What is einsum_path()?

The einsum_path() function evaluates different contraction orders for Einstein summation operations. It returns a tuple containing the optimal contraction path and detailed information about the optimization process.

Syntax

numpy.einsum_path(subscripts, *operands, optimize='greedy')

Parameters

• subscripts ? String specifying the subscripts for summation using Einstein notation
• operands ? Input arrays for the operation
• optimize ? Optimization strategy ('greedy', 'optimal', or False)

Basic Example

Let's find the optimal contraction order for a matrix chain multiplication ?

import numpy as np

# Set seed for reproducible results
np.random.seed(123)

# Create tensors
p = np.random.rand(2, 2)
q = np.random.rand(2, 5) 
r = np.random.rand(5, 2)

# Find optimal contraction path
path_info = np.einsum_path('ij,jk,kl->il', p, q, r, optimize='greedy')

print("Optimal path:", path_info[0])
print("\nDetailed information:")
print(path_info[1])

Optimal path: ['einsum_path', (1, 2), (0, 1)]

Detailed information:
  Complete contraction:  ij,jk,kl->il
         Naive scaling:  4
     Optimized scaling:  3
      Naive FLOP count:  1.200e+02
  Optimized FLOP count:  5.700e+01
   Theoretical speedup:  2.105
  Largest intermediate:  4.000e+00 elements
--------------------------------------------------------------------------
scaling                 current                     remaining
--------------------------------------------------------------------------
   3           kl,jk->jl                  ij,jl->il
   3                 jl,ij->il                     il->il

Understanding the Output

The path [(1, 2), (0, 1)] means ?

• First, contract operands 1 and 2 (q and r): jk,kl->jl
• Then contract the result with operand 0 (p): ij,jl->il

Comparison of Optimization Strategies

import numpy as np

np.random.seed(123)
a = np.random.rand(10, 10)
b = np.random.rand(10, 20)
c = np.random.rand(20, 5)

# Compare different optimization strategies
strategies = ['greedy', 'optimal']

for strategy in strategies:
    path_info = np.einsum_path('ij,jk,kl->il', a, b, c, optimize=strategy)
    print(f"{strategy.capitalize()} strategy:")
    print(f"Path: {path_info[0]}")
    print(f"FLOP count: {path_info[1].split('Optimized FLOP count:')[1].split()[0]}")
    print()

Greedy strategy:
Path: ['einsum_path', (0, 1), (0, 1)]
FLOP count: 3.000e+03

Optimal strategy:
Path: ['einsum_path', (1, 2), (0, 1)]
FLOP count: 3.000e+03

Key Benefits

Conclusion

Use numpy.einsum_path() to find the optimal contraction order for complex tensor operations. This optimization can significantly reduce computational cost and memory usage, especially for large tensor contractions.