This package lets you work with multi-dimensional arrays in index notation, by defining a few macros which translate this to broadcasting, permuting, and reducing operations.
The first is
@cast, which deals both with "casting" into new shapes (including going to and from an array-of-arrays) and with broadcasting:
@cast A[row][col] := B[row, col] # slice a matrix B into rows, also @cast A[r] := B[r,:] @cast C[(i,j), (k,ℓ)] := D.x[i,j,k,ℓ] # reshape a 4-tensor D.x to give a matrix @cast E[φ,γ] = F[φ]^2 * exp(G[γ]) # broadcast E .= F.^2 .* exp.(G') into existing E @cast _[i] := isodd(i) ? log(i) : V[i] # broadcast a function of the index values @cast T[x,y,n] := outer(M[:,n])[x,y] # generalised mapslices, vector -> matrix function
@reduce takes sums (or other reductions) over the indicated directions. Among such sums is
matrix multiplication, which can be done more efficiently using
@reduce K[_,b] := prod(a,c) L.field[a,b,c] # product over dims=(1,3), drop dims=3 @reduce S[i] = sum(n) -P[i,n] * log(P[i,n]/Q[n]) # sum!(S, @. -P*log(P/Q')) into exising S @matmul M[i,j] := sum(k,k′) U[i,k,k′] * V[(k,k′),j] # matrix multiplication, plus reshape
The same notation with
@cast applies a function accepting the
dims keyword, without reducing:
@cast W[i,j,c,n] := cumsum(c) X[c,i,j,n]^2 # permute, broadcast, cumsum(; dims=3)
All of these are converted into array commands like
eachslice, plus a broadcasting expression if needed,
mul!. This package just provides a convenient notation.
using OffsetArrays @cast R[n,c] := n^2 + rand(3)[c] (n in -5:5) # arbitrary indexing
And it can be used with some packages which modify broadcasting:
using Strided, LoopVectorization, LazyArrays @cast @strided E[φ,γ] = F[φ]^2 * exp(G[γ]) # multi-threaded @reduce @avx S[i] := sum(n) -P[i,n] * log(P[i,n]) # SIMD-enhanced @reduce @lazy M[i,j] := sum(k) U[i,k] * V[j,k] # non-materialised
using Pkg; Pkg.add("TensorCast")
Similar notation is also used by some other packages, although all of them use an implicit sum over repeated indices. TensorOperations.jl performs Einstein-convention contractions and traces:
@tensor A[i] := B[i,j] * C[j,k] * D[k] # matrix multiplication, A = B * C * D @tensor D[i] := 2 * E[i] + F[i,k,k] # partial trace of F only, Dᵢ = 2Eᵢ + Σⱼ Fᵢⱼⱼ
More general contractions are allowed by OMEinsum.jl, but only one term:
@ein Z[i,j,ξ] := X[i,k,ξ] * Y[j,k,ξ] # batched matrix multiplication Z = ein" ikξ,jkξ -> ijξ "(X,Y) # numpy-style notation
@einsum S[i] := -P[i,n] * log(P[i,n]/Q[n]) # sum over n, for each i (also with @reduce above) @einsum G[i] := 2 * E[i] + F[i,k,k] # the sum includes everyting: Gᵢ = Σⱼ (2Eᵢ + Fᵢⱼⱼ) @tullio Z[i,j] := abs(A[i+x, j+y] * K[x,y]) # convolution, summing over x and y
These produce very different code for actually doing what you request:
@ein work out a sequence of basic operations (like contraction and traces),
@tullio write the necessary set of nested loops directly (plus optimisations).
This was a holiday project to learn a bit of metaprogramming, originally
But it suffered a little scope creep.