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

To get the lowest cost contraction order for an einsum expression, use the numpy.einsum+path() method in Python. The 1st parameter, subscripts specify the subscripts for summation. The 2nd parameter, operands are the arrays for the operation.

Using the Einstein summation convention, many common multi-dimensional, linear algebraic array operations can be represented in a simple fashion. In implicit mode einsum computes these values.

In explicit mode, einsum provides further flexibility to compute other array operations that might not be considered classical Einstein summation operations, by disabling, or forcing summation over specified subscript labels.

The resulting path indicates which terms of the input contraction should be contracted first, the result of this contraction is then appended to the end of the contraction list. This list can then be iterated over until all intermediate contractions are complete.

Steps

At first, import the required libraries −

import numpy as np

Tensors −

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

To get the lowest cost contraction order for an einsum expression, use the numpy.einsum+path() method −

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

Displaying Path info −

print(path_info[0])
print(path_info[1])

Example

import numpy as np
np.random.seed(123)

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

# To get the lowest cost contraction order for an einsum expression, use the numpy.einsum+path() method in Python
path_info = np.einsum_path('ij,jk,kl->il', p, q, r, optimize='greedy')

# Displaying Path info
print(path_info[0])
print(path_info[1])

Output

['einsum_path', (1, 2), (0, 1)]
  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