Contractions efficaces avec les chemins einsum de NumPy

Question

J'ai un ensemble de contractions que je voudrais optimiser ; pour les contractions, j'utilise np.einsum() du module NumPy. L'exemple minimal reproductible est ici :

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)

Les temps semblent se situer autour de ~2 secondes. J'ai également observé que einsum n'est généralement pas multithreadé, et un seul cœur est utilisé à la place. Existe-t-il un autre moyen efficace d'effectuer de telles contractions ? (Peut-être avec Numba ?).

Answer 1

Au cas où cela aiderait à la discussion, voici le premier. 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

Par simple édition de chaînes, j'ai retravaillé le path1 sur

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

abefi" de a ne font que passer ; de même, "mnqrt" pour la dernière étape de l'opération. a ; "i" et "t" sont (100,) les autres sont petits.

cg', 'dh', 'ko', 'lp' sont des contractions entre paires (celles-ci sont petites).

"j" et "s" apparaissent dans 3 tableaux. (les deux sont (100,))

Avec (comme suggéré dans le commentaire) :

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) effectue un calcul par élément diffusé sans aucune somme de produits. js' est la grande paire de dimensions (100,100). L'ensemble de contraction est vide.

(2,0) fait la contraction sur {'d', 'j', 'g', 'c', 'h'} dont j est le plus grand.

(1,0) complète la contraction sur {'l', 'o', 'k', 's', 'p'}