numpy.einsum per Julia? (2)

Question

Proveniente da questo question, mi chiedo se sia possibile un einsum più generalizzato. Supponiamo, ho avuto il problema

using PyCall 
@pyimport numpy as np 

a = rand(10,10,10) 
b = rand(10,10) 
c = rand(10,10,10) 

Q = np.einsum("imk,ml,lkj->ij", a,b,c) 

o qualcosa di simile, come se fossi a risolvere questo problema senza scorrendo le somme?

con i migliori saluti

Answer 1

Modifica/Aggiornamento: questo è ora un pacchetto registrato, in modo da poter Pkg.add("Einsum") e si dovrebbe essere pronti per partire (vedere l'esempio di seguito per iniziare).

Risposta originale: Ho appena creato un codice molto preliminare per farlo. Segue esattamente quello che Matt B. ha descritto nel suo commento. Spero che aiuti, fammi sapere se ci sono problemi con esso.

https://github.com/ahwillia/Einsum.jl

Questo è come si dovrebbe implementare il tuo esempio:

using Einsum 

a = rand(10,10,10) 
b = rand(10,10) 
c = rand(10,10,10) 
Q = zeros(10,10) 

@einsum Q[i,j] = a[i,m,k]*b[m,l]*c[l,k,j] 

Sotto il cofano della macro costruisce la seguente serie di annidati cicli for e li inserisce nel codice prima del tempo di compilazione. (Nota: questo non è il codice esatto inserito, esso controlla anche per assicurarsi che le dimensioni degli ingressi sono d'accordo, utilizzando macroexpand per vedere il codice completo):

for j = 1:size(Q,2) 
    for i = 1:size(Q,1) 
     s = 0 
     for l = 1:size(b,2) 
      for k = 1:size(a,3) 
       for m = 1:size(a,2) 
        s += a[i,m,k] * b[m,l] * c[l,k,j] 
       end 
      end 
     end 
     Q[i,j] = s 
    end 
end