5 Best Ways to Evaluate the Lowest Cost Contraction Order for an einsum Expression in Python

by

💡 Problem Formulation: Evaluating tensor expressions using Einsum (Einstein summation convention) in Python can become computationally intensive, especially for large tensors with complex operations. An optimal contraction order can significantly reduce computation time and resources. This article discusses strategies to identify the lowest cost contraction path for an Einsum expression. For example, given the expression 'ijk,jl->il', with the corresponding input tensors of appropriate dimensions, we aim to find the most efficient computation sequence.

Method 1: Naive Evaluation

The naive method of contraction order evaluation involves attempting every permutation of contraction orders and selecting the one with minimal cost. This is akin to brute-forcing and often impractical for larger expressions due to exponential growth in possibilities. The method uses the numpy.einsum_path() function with its ‘optimal’ strategy, which tries to find the best contraction order but can be slow for large expressions.

Here’s an example:

import numpy as np

# Define the tensors
A = np.random.rand(10, 10, 10)
B = np.random.rand(10, 10)

# Evaluate the einsum path
path_info = np.einsum_path('ijk,jl->il', A, B, optimize='optimal')
print(path_info[1])

The output could look like this:

Einsum path:
  -> ijk,jl->il
     (ij|k),jl->il
       ijl->il

This code snippet imports the numpy library, generates two random tensors A and B, and applies the einsum_path() function with the ‘optimal’ strategy to get the best contraction order. The output describes the sequence of operations for the einsum expression.

Method 2: Greedy Algorithm

The greedy algorithm for finding a contraction order is a heuristic that at each step, selects the local optimum in the hopes of finding a global optimum. In the context of einsum expressions, the greedy algorithm aims to minimize intermediate tensor sizes. This method is implemented in the np.einsum_path() function with the ‘greedy’ strategy available.

Here’s an example:

path_info = np.einsum_path('ijk,jl->il', A, B, optimize='greedy')
print(path_info[1])

The output might be:

Einsum path:
  -> ijk,jl->il
     ijkl->il
       il->il

By requesting a ‘greedy’ optimization, this example finds an efficient contraction path that doesn’t necessarily guarantee the optimal solution but provides a good balance between performance and evaluation time.

Method 3: Dynamic Programming

Dynamic programming is an optimization technique that solves complex problems by breaking them down into simpler subproblems. It is a more sophisticated approach than the greedy algorithm and often more efficient than the naive method. While not directly provided by the NumPy einsum function, external libraries like opt_einsum can apply dynamic programming to find contraction paths.

Here’s an example:

from opt_einsum import contract_path

# Path optimization using dynamic programming
path_info_dp = contract_path('ijk,jl->il', A, B, optimize='dp')
print(path_info_dp[1])

The output could be:

Einsum path:
  -> ijk,jl->il
     ijkl->il
       il->il

This code example solves the einsum expression using dynamic programming to achieve an efficient contraction order. It leverages the opt_einsum library’s contract_path() function with the ‘dp’ optimization strategy.

Method 4: Custom Heuristic

Developing a custom heuristic to find an einsum contraction path involves crafting your own set of rules and criteria to approximate the best order of operations. This could be useful when domain-specific knowledge can guide the contraction sequence. While potentially outperforming generic algorithms, creating a custom heuristic is usually a complex and time-consuming process.

Here’s an example:

# A mock-up function illustrating a custom heuristic approach
def custom_heuristic_einsum_path(expression, *tensors):
    # Custom logic to determine the contraction order
    contraction_order = 'custom logic here'
    return contraction_order

# Use the custom heuristic function
path = custom_heuristic_einsum_path('ijk,jl->il', A, B)
print(path)

The output would be a representation of the custom heuristic’s contraction order:

'custom logic here'

This code snippet pretends to implement a function for computing the einsum path using a custom heuristic. The function, while not operational, is meant to highlight how such a method could be structured.

Bonus One-Liner Method 5: Default

The default method is the simplest approach which relies on NumPy’s built-in optimization strategy when none is specified. By default, numpy.einsum chooses an optimization strategy between ‘optimal’ and ‘greedy’ based on the expression size.

Here’s an example:

path_info_default = np.einsum_path('ijk,jl->il', A, B)
print(path_info_default[1])

The output will resemble:

Einsum path:
  -> ijk,jl->il
     ijkl->il
       il->il

By calling einsum_path() without an explicit ‘optimize’ argument, NumPy determines an appropriate contraction strategy. This is a suitable default choice for many cases due to its simplicity.

Summary/Discussion

  • Method 1: Naive Evaluation. The ‘optimal’ strategy is extremely thorough, making it suitable for small expressions but impractical for large-scale problems due to its high computational cost.
  • Method 2: Greedy Algorithm. The ‘greedy’ strategy strikes a balance between execution time and evaluation quality, making it a robust choice for medium to large expressions.
  • Method 3: Dynamic Programming. Although it requires an external library, dynamic programming stands out for larger expressions by providing near-optimal solutions in reasonable time frames.
  • Method 4: Custom Heuristic. Custom heuristics are powerful when tailor-made to the application but often require extensive domain knowledge and can be difficult to generalize.
  • Method 5: Default. The default optimization in NumPy is an excellent starting point, as it simplifies the process and provides satisfactory results for most applications without additional complexity.