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

In this article, we will see how to Evaluate the lower cost contraction for an einsum expression in Python.

np.einsum_path method

To obtain the lowest cost contraction order for an einsum expression in Python, use the numpy.einsum path() function. The Einstein summation convention can be used to represent several common multi-dimensional linear algebraic array operations. In implicit mode, Einsum calculates these values. Einsum gives additional flexibility to compute alternative array operations that may not be considered standard Einstein summation operations in explicit mode by disabling or forcing summing across defined subscript labels. The path indicates which terms should be contracted first, with the result appended to the end of the contraction list. This list can then be iterated until all intervening contractions have been completed.

Syntax: np.einsum_path(subscripts, *operands, optimize=’greedy’):

Parameters:

• subscripts: string type .The subscripts for summation are specified.
• *operands : list of array_like objects.These are the operation’s arrays.
• optimize : {bool, list, tuple, ‘greedy’, ‘optimal’}. Select the type of path.

Returns:

• path : list of tuples.The einsum path is represented as a list.
• string_repr: str. The einsum path as a printable representation.

Example

The libraries are imported here, and the np.random.rand() method is used to generate random number arrays. We make three arrays called a, b, and c. The lower cost contraction for an einsum expression is evaluated using the np.einsum path() method. Let’s begin with a chain dot illustration. The b and c tensors should be contracted first in this situation, as shown by the path’s first element (1, 2). The final contraction (0, 1) is completed by appending the resulting tensor to the end of the contraction.

Python3

# importing packages
import numpy as np

# setting seed
np.random.seed(101)

# creating random arrays
a = np.random.rand(2, 2)
print(a)
b = np.random.rand(2, 3)
print(b)
c = np.random.rand(3, 2)
print(c)

# lowest cost contraction order for an einsum expression
path = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy')

# Path info
print(path[0])
print(path[1])

Output:

[[0.51639863 0.57066759]

 [0.02847423 0.17152166]]

[[0.68527698 0.83389686 0.30696622]

 [0.89361308 0.72154386 0.18993895]]

[[0.55422759 0.35213195]

 [0.1818924  0.78560176]

 [0.96548322 0.23235366]]

[‘einsum_path’, (1, 2), (0, 1)]

  Complete contraction:  ij,jk,kl->il

         Naive scaling:  4

     Optimized scaling:  3

      Naive FLOP count:  7.200e+01

  Optimized FLOP count:  4.100e+01

   Theoretical speedup:  1.756

  Largest intermediate:  4.000e+00 elements

————————————————————————–

scaling                  current                                remaining

————————————————————————–

   3                   kl,jk->jl                                ij,jl->il

   3                   jl,ij->il                                   il->il