Einstein summation notation similar to numpy's `einsum`

function (but more flexible!).

PackageEvaluator |
Package Build |
Package Status |
---|---|---|

- help wanted! |

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

.

`=`

```
using Einsum
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]
```

`:=`

```
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] # creates new array A with appropriate dimensions
```

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

To see exactly what is generated, use `macroexpand`

:

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

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

Again, you can use `macroexpand`

to see the exact code that is generated.

04/16/2016

2 months ago

71 commits