Efficient contractions with NumPy einsum paths

user_2261553 0
Viewed 99

I have a set of contractions that I would like to optimize; for the contractions I am using np.einsum() from the NumPy module. The minimal reproducible example is here:

import numpy as np
from time import time

d1=2
d2=3
d3=100

a = np.random.rand( d1,d1,d1,d1,d2,d2,d2,d2,d3,d3 ) + 1j*np.random.rand( d1,d1,d1,d1,d2,d2,d2,d2,d3,d3 )
b = np.random.rand( d1,d1, d2,d2,d3,d3 ) + 1j*np.random.rand( d1,d1,d2,d2,d3,d3 )
c = np.random.rand( d1,d1, d2,d2,d3,d3 ) + 1j*np.random.rand( d1,d1,d2,d2,d3,d3 )

path_1  = 'abcdefghij,ckgojs,dlhpjs,klmnopqrst->abmnefqrit'
path_2  = 'abcdefghij,ckgoji,nbrfji,klmnopqrij->almdepqhij'

ts = time()
einsum_pathinfo = np.einsum_path(path_1, a, b, c, a )
term_a          = np.einsum(path_1, a, b, c, a, optimize=einsum_pathinfo[0])
print("took", time()-ts)

ts = time()
einsum_pathinfo = np.einsum_path( path_2, a, b, c , a )
term_a          = np.einsum(path_2, a, b, c, a, optimize=einsum_pathinfo[0])
print("took", time()-ts)

The times seem to be around ~2 seconds. I have also observed that einsum is generally not multithreaded, and a single core is used instead. Is there any other efficient way to perform such contractions? (Perhaps with Numba?).

numpy-einsum python-3.x numpy
1 Answers

In case it helps with the discussion, here's the first pathinfo:

In [241]: print(einsum_pathinfo[0])
['einsum_path', (1, 2), (0, 2), (0, 1)]

In [242]: print(einsum_pathinfo[1])
  Complete contraction:  abcdefghij,ckgojs,dlhpjs,klmnopqrst->abmnefqrit
         Naive scaling:  20
     Optimized scaling:  15
      Naive FLOP count:  6.718e+14
  Optimized FLOP count:  1.866e+11
   Theoretical speedup:  3599.750
  Largest intermediate:  1.296e+07 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
  10    dlhpjs,ckgojs->cdklghopjs abcdefghij,klmnopqrst,cdklghopjs->abmnefqrit
  15    cdklghopjs,abcdefghij->abklefopis        klmnopqrst,abklefopis->abmnefqrit
  15    abklefopis,klmnopqrst->abmnefqrit                   abmnefqrit->abmnefqrit

Just by string editing I reworked the path1 into

'(abefi)(cg)(dh)j,(cg)(ko)js,(dh)(lp)js,(ko)(lp)(mnqrt)s->(abefi)(mnqrt)'

'abefi' from a just pass through; likewise 'mnqrt' for the last a; 'i' and 't' are (100,) the others are small.

'cg', 'dh', 'ko', 'lp' are contractions between pairs (these are small)

'j' and 's' appear in 3 arrays. (both of these are (100,))

With (as suggested in the comment):

einsum_pathinfo = np.einsum_path(path_1, a, b, c, a ,einsum_call=True)

In [249]: print(einsum_pathinfo[1])
[((2, 1), set(), 'dlhpjs,ckgojs->cdklghopjs', 
 ['abcdefghij', 'klmnopqrst', 'cdklghopjs'], 
 False), 
 ((2, 0), {'d', 'j', 'g', 'c', 'h'}, 'cdklghopjs,abcdefghij->abklefopis', 
 ['klmnopqrst', 'abklefopis'], 
 True), 
 ((1, 0), {'l', 'o', 'k', 's', 'p'}, 'abklefopis,klmnopqrst->abmnefqrit', 
 ['abmnefqrit'], 
 True)]

(2,1) is doing a broadcasted element-wise calculation without any sum-of-products. 'js' is the large (100,100) dimensions pair. The contraction set is empty.

(2,0) does the contraction on {'d', 'j', 'g', 'c', 'h'}, of which j is largest.

(1,0) completes the contraction on {'l', 'o', 'k', 's', 'p'}

user_9019250
Related