TensorCast.jl

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

Second, @reduce takes sums (or other reductions) over the indicated directions. Among such sums is matrix multiplication, which can be done more efficiently using @matmul instead:

@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 reshape and permutedims and eachslice, plus a broadcasting expression if needed, and sum / sum!, or * / mul!. This package just provides a convenient notation.

[!WARNING] Writing @reduce C[i,j] := sum(k) A[i,k] * B[k,j] is terrible way to perform matrix multiplication. This creates a huge array A .* reshape(B, 1, size(B)...) before summing, which is much slower than the built-in A * B. See below for other packages which aim to be good at such operations.

From version 0.4, it relies on TransmuteDims.jl to handle re-ordering of dimensions, and LazyStack.jl to handle slices. It should also now work with OffsetArrays.jl:

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 @turbo 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

It should work automatically with most array types. This includes GPU arrays such as those from CUDA.jl, whose broadcasting is executed on the device.

Installation

using Pkg; Pkg.add("TensorCast")

The current version requires Julia 1.6 or later. There are a few pages of documentation.

Elsewhere

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.jl and Tullio.jl allow arbitrary (element-wise) functions:

@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

Notice that @einsum and @tullio sum the entire right hand side, like @reduce does, while @tensor sums individual terms.

These produce very different code for actually doing what you request: The macros @tensor and @ein work out a sequence of basic tensor operations (like contraction and traces), while @einsum and @tullio write the necessary set of nested loops directly (plus optimisations). This package's macros @cast, @reduce and @matmul instead write everything in terms of whole-array operations (like reshape, permutedims and broadcasting).

For those who speak Python, @cast and @reduce allow similar operations to einshape or einops (minus the cool video, but plus broadcasting). In the tests, this file translates many examples. Python's einsum is closer OMEinsum.@ein and TensorOperations.@tensor, and this package's @matmul.

About

This was a holiday project to learn a bit of metaprogramming, originally TensorSlice.jl. But it suffered a little scope creep.