# ahwillia/Einsum.jl

Einstein summation notation in Julia
Switch branches/tags
Nothing to show
Fetching latest commit…
Cannot retrieve the latest commit at this time.
 Failed to load latest commit information. perf Jul 7, 2016 src Aug 10, 2018 test Apr 22, 2018 .travis.yml Apr 7, 2018 LICENSE.md Sep 4, 2016 README.md Apr 22, 2018 REQUIRE Apr 22, 2018

# Einsum.jl

Einstein summation notation similar to numpy's `einsum` function (but more flexible!).

PackageEvaluator Package Build Package Status
- help wanted!

To install: `Pkg.add("Einsum")`.

## Documentation

### Basics

If the destination array is preallocated, then use `=`:

```A = zeros(5,6,7) # need to preallocate destination
X = randn(5,2)
Y = randn(6,2)
Z = randn(7,2)
@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 A with appropriate dimensions:

```using Einsum
X = randn(5,2)
Y = randn(6,2)
Z = randn(7,2)
@einsum A[i,j,k] := X[i,r]*Y[j,r]*Z[k,r]```

### What happens under the hood

To see exactly what is generated, use `@macroexpand`:

`@macroexpand @einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]`

The `@einsum` macro automatically generates code that looks much like the following (note that we "sum out" over the index `r`, since it only occurs on the right hand side of the equation):

```for k = 1:size(A,3)
for j = 1:size(A,2)
for i = 1:size(A,1)
s = 0
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```

In reality, this code will be preceded by the the neccessary bounds checks and allocations, and take care to use the right types and keep hygenic.

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
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```

### `@einsimd`

This is a variant of `@einsum` which will put `@simd` in front of the innermost loop; e.g., `@einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]` will result approximately in

```for k = 1:size(A,3)
for j = 1:size(A,2)
for i = 1:size(A,1)
s = 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

In principle, the `@einsum` macro can flexibly 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.

### Related Packages:

• TensorOperations.jl has less flexible syntax (and does not allow certain 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 (cache-friendly) order for the looping.