Einsum.jl
Package Build  Package Status 

 help wanted! 
This package exports a single macro @einsum
, which implements similar notation to the Einstein
summation convention to flexibly specify
operations on Julia Array
s, similar to numpy's
einsum
function
(but more flexible!).
For example, basic matrix multiplication can be implemented as:
@einsum A[i, j] := B[i, k] * C[k, j]
To install: Pkg.add("Einsum")
, or else pkg> add Einsum
after pressing ]
on Julia 0.7 and later.
Documentation
Basics
If the destination array is preallocated, then use =
:
A = ones(5, 6, 7) # preallocated space
X = randn(5, 2)
Y = randn(6, 2)
Z = randn(7, 2)
# Store the result in A, overwriting as necessary:
@einsum A[i, j, k] = X[i, r] * Y[j, r] * Z[k, r]
If destination is not preallocated, then use :=
to automatically create a new array for the result:
X = randn(5, 2)
Y = randn(6, 2)
Z = randn(7, 2)
# Create new array B with appropriate dimensions:
@einsum B[i, j, k] := X[i, r] * Y[j, r] * Z[k, r]
What happens under the hood
To execute an expression, @einsum
uses Julia's metaprogramming
capabilities to generate and execute a
series of nested for loops. To see exactly what is generated, use
@macroexpand
(or @expand
from MacroTools.jl):
@macroexpand @einsum A[i, j] := B[i, k] * C[k, j]
The output will look much like the following (note that we "sum out" over the index k
, since it
only occurs multiple times on the right hand side of the equation):
# determine type
T = promote_type(eltype(B), eltype(C))
# allocate new array
A = Array{T}(undef, size(B))
# check dimensions
@assert size(B, 2) == size(C, 2)
# main loop
@inbounds begin # skip boundschecking for speed
for i = 1:size(B, 1), j = 1:size(C, 2)
s = zero(T)
for k = 1:size(B,2)
s += B[i, k] * C[k, j]
end
A[i, j] = s
end
end
The actual generated code is a bit more verbose (and not neatly commented/formatted), and will take care to use the right types and keep hygienic.
You can also use updating assignment operators for preallocated arrays. E.g., @einsum A[i, j, k] *= X[i, r] * Y[j, r] * Z[k, r]
will produce something like
for k = 1:size(A, 3)
for j = 1:size(A, 2)
for i = 1:size(A, 1)
s = 0.0
for r = 1:size(X, 2)
s += X[i, r] * Y[j, r] * Z[k, r]
end
# Difference: here, the updating form is used:
A[i, j, k] *= s
end
end
end
Rules for indexing variables
 Indices that show up on the righthandside but not the lefthandside are summed over
 Arrays which share an index must be of the same size, in those dimensions
@einsum
iterates over the extent of the righthandside indices. For example, the following code
allocates an array A
that is the same size as B
and copies its data into A
:
@einsum A[i, j] := B[i, j] # same as A = copy(B)
If an index appears on the righthandside, but does not appear on the lefthandside, then this
variable is summed over. For example, the following code allocates A
to be size(B, 1)
and sums
over the rows of B
:
@einsum A[i] := B[i, j] # same as A = dropdims(sum(B, dims=2), dims=2)
If an index variable appears multiple times on the righthandside, then it is asserted that the sizes of these dimensions match. For example,
@einsum A[i] := B[i, j] * C[j]
will check that the second dimension of B
matches the first dimension of C
in length. In
particular it is equivalent to the following code:
A = zeros(size(B, 1))
@assert size(B, 2) == size(C, 1)
for i = 1:size(B, 1), j = 1:size(B, 2)
A[i] += B[i, j] * C[j]
end
So an error will be thrown if the specified dimensions of B
and C
don't match.
Offset indexing
@einsum
also allows offsets on the righthandside, the range over which i
is summed is then restricted:
@einsum A[i] = B[i  5]
@vielsum
This variant of @einsum
will run multithreaded on the outermost loop. For this to be fast, the code must not introduce temporaries like s = 0
in the example above. Thus for example @expand @vielsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]
results in something equivalent to @expand
ing the following:
Threads.@threads for k = 1:size(A,3)
for j = 1:size(A,2)
for i = 1:size(A,1)
A[i,j,k] = 0.0
for r = 1:size(X,2)
A[i,j,k] += X[i,r] * Y[j,r] * Z[k,r]
end
end
end
end
For this to be useful, you will need to set an environment variable before starting Julia, such as export JULIA_NUM_THREADS=4
. See the manual for details, and note that this is somewhat experimental. This will not always be faster, especially for small arrays, as there is some overhead to dividing up the work.
At present you cannot use updating assignment operators like +=
with this macro (only =
or :=
) and you cannot assign to a scalar lefthandside (only an array). These limitations prevent different threads from writing to the same memory at the same time.
@einsimd
This is a variant of @einsum
which will put @simd
in front of the innermost loop, encouraging Julia's compiler vectorize this loop (although it may do so already). For example @einsimd A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]
will result in approximately
@inbounds for k = 1:size(A,3)
for j = 1:size(A,2)
for i = 1:size(A,1)
s = 0.0
@simd for r = 1:size(X,2)
s += X[i,r] * Y[j,r] * Z[k,r]
end
A[i,j,k] = s
end
end
end
Whether this is a good idea or not you have to decide and benchmark for yourself in every specific case. @simd
makes sense for certain kinds of operations; make yourself familiar with its documentation and the inner workings of it in general.
Other functionality
The @einsum
macro can implement function calls within the nested for loop structure. For example, consider transposing a block matrix:
z = Any[rand(2,2) for i=1:2, j=1:2]
@einsum t[i,j] := transpose(z[j,i])
This produces a for loop structure with a transpose
function call in the middle. Approximately:
for j = 1:size(z,1)
for i = 1:size(z,2)
t[i,j] = transpose(z[j,i])
end
end
This will work as long the function calls are outside the array names. Again, you can use @macroexpand
to see the exact code that is generated.
The output need not be an array. But note that on Julia 0.7 and 1.0, the rules for evaluating in global scope (for example at the REPL prompt) are a little different  see this package for instance (which is loaded in IJulia notebooks). To get the same behavior as you would have inside a function, you use a let
block:
p = rand(5); p .= p ./ sum(p)
let
global S
@einsum S :=  p[i] * log(p[i])
end
Or perhaps clearer, explicitly define a function:
m(pq, p, q) = @einsum I := pq[i,j] * log(pq[i,j] / p[i] / q[j])
m(rand(5,10), rand(5), rand(10))
Related Packages:

TensorOperations.jl has less flexible syntax (only allowing strict Einstein convention contractions), but can produce much more efficient code. Instead of generating “naive” loops, it transforms the expressions into optimized contraction functions and takes care to use a good (cachefriendly) order for the looping.

ArrayMeta.jl aims to produce cachefriendly operations for more general loops (but is Julia 0.6 only).