numpy.einsum for Julia? (2)

Question

Coming from this question, I wonder if a more generalized einsum was possible. Let us assume, I had the 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 were I to solve this problem without looping through the sums?

with best regards

Answer 1

Edit/Update: This is now 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 to do this. It follows exactly what Matt B. described in his comment. Hope it helps, let me know if there are problems with it.

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

This is how you would 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 the macro builds the following series of nested for loops and inserts them into your code before compile time. (Note this is not the exact code inserted, it also checks to make sure the dimensions of the inputs 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