TensorOperations.jl

Fast tensor operations using a convenient Einstein index notation.

Table of contents

Installation

Install with the package manager, pkg> add TensorOperations.

Package features

The TensorOperations.jl package is centered around the following features:

  • A macro @tensor for conveniently specifying tensor contractions and index permutations via Einstein's index notation convention. The index notation is analyzed at compile time and lowered into primitive tensor operations, namely (permuted) linear combinations and inner and outer contractions. The macro supports several keyword arguments to customize the lowering process, namely to insert additional checks that help with debugging, to specify contraction order or to automatically determine optimal contraction order for given costs (see next bullet), and finally, to select different backends to evaluate those primitive operations.
  • The ability to optimize pairwise contraction order in complicated tensor contraction networks according to the algorithm in this paper, where custom (compile time) costs can be specified, either as a keyword to @tensor or using the @tensoropt macro (for expliciteness and backward compatibility). This optimization is performed at compile time, and the resulting contraction order is hard coded into the resulting expression. The similar macro @tensoropt_verbose provides more information on the optimization process.
  • A function ncon (for network contractor) for contracting a group of tensors (a.k.a. a tensor network), as well as a corresponding @ncon macro that simplifies and optimizes this slightly. Unlike the previous macros, ncon and @ncon do not analyze the contractions at compile time, thus allowing them to deal with dynamic networks or index specifications.
  • (Experimental) support for automatic differentiation by supplying chain rules for the different tensor operations using the ChainRules.jl interface.
  • The ability to support different tensor types by overloading a minimal interface of tensor operations, or to support different implementation backends for the same tensor type.
  • An efficient default implementation for Julia Base arrays that qualify as strided, i.e. such that its entries are layed out according to a regular pattern in memory. The only exceptions are ReinterpretedArray objects. Additionally, Diagonal objects whose underlying diagonal data is stored as a strided vector are supported. This facilitates tensor contractions where one of the operands is e.g. a diagonal matrix of singular values or eigenvalues, which are returned as a Vector by Julia's eigen or svd method. This implementation for AbstractArray objects is based on Strided.jl for efficient (cache-friendly and multithreaded) tensor permutations (transpositions) and gemm from BLAS for contractions. There is also a fallback contraction strategy that is natively built using Strided.jl, e.g. for scalar types which are not supported by BLAS. Additional backends (e.g. pure Julia Base using loops and/or broadcasting) may be added in the future.
  • Support for CuArray objects if used together with CUDA.jl and cuTENSOR.jl, by relying on (and thus providing a high level interface into) NVidia's cuTENSOR library.

Tensor operations

TensorOperations.jl supports 3 basic tensor operations, i.e. primitives in which every more complicated tensor expression is deconstructed.

  1. addition: Add a (possibly scaled version of) one tensor to another tensor, where the indices of both arrays might appear in different orders. This operation combines normal tensor addition (or linear combination more generally) and index permutation. It includes as a special case copying one tensor into another with permuted indices.
  2. trace or inner contraction: Perform a trace/contraction over pairs of indices of a single tensor array, where the result is a lower-dimensional array.
  3. (outer) contraction: Perform a general contraction of two tensors, where some indices of one array are paired with corresponding indices in a second array.

To do list

  • Add more backends, e.g. using pure Julia Base functionality, or using LoopVectorization.jl
  • Make it easier to modify the contraction order algorithm or its cost function (e.g. to optimize based on memory footprint) or to splice in runtime information.