Numpy.einsum for Julia? (2)

Question

Based on this question , I wonder if a more generalized einsum is possible. Suppose I had a problem

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) 

Or something similar, how could I solve this problem without sorting through the sums?

with best wishes

Answer 1

Modify / Update: Now this is a registered package, so you can Pkg.add("Einsum") and you should be good to go (see the example below to get started).

Original answer: I just created some very preliminary code for this. This is exactly what Matt B. described in his commentary. Hope this helps, let me know if there are problems with this.

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

Here's how you could implement your example:

 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] 

Under the hood, a macro builds the following sequence of nested loops and inserts them into your code until compilation. (Note that this is not an exact code entered, it also checks to see if the input sizes agree using macroexpand to see the full code):

 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