TensorOperations.jl writes tensor contractions in Einstein notation

What the @tensor macro actually does

TensorOperations.jl is mostly used through the @tensor macro, which lets you express a given operation in index notation, also known as Einstein notation, using the Einstein summation convention. The description of the package is fast tensor operations using a convenient Einstein index notation.

The mechanics are worth spelling out because they are the whole value proposition. In Einstein notation, an index that appears once on each side of an equation is summed over implicitly, and an index that appears twice on the right and once on the left is a contraction index. The macro reads that pattern out of the expression and generates the loop nest that computes it, so you write the mathematics rather than the iteration.

The readme example makes the two assignment forms concrete:

using TensorOperations
α = randn()
A = randn(5, 5, 5, 5, 5, 5)
B = randn(5, 5, 5)
C = randn(5, 5, 5)
D = zeros(5, 5, 5)
@tensor begin
    D[a, b, c] = A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b]
    E[a, b, c] := A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b]
end

The difference between the last two lines is the difference between writing into a preallocated array and allocating a new one. In the second to last line the result is stored in the preallocated array D, whereas the last line uses a different assignment operator, :=, in order to define and allocate a new array E of the correct size. The contents of D and E will be equal.

That two-character distinction is the API's main ergonomic decision, and it is the one you will use most often.

Backends and allocation strategies are pluggable in v5

The v5 release notes describe three changes and one warning, and together they say a lot about the direction of the project.

First, there is support for cuTENSOR v2, and with that more recent versions of CUDA.jl. cuTENSOR is NVIDIA's tensor contraction library, and supporting it means the GPU path is not restricted to what a hand-written Julia kernel can do.

Second, there is improved support for automatic differentiation using reverse-mode rules with ChainRulesCore.jl, which is what makes the package usable inside a differentiable pipeline rather than only in a forward computation.

Third, there is improved and extended support for backends and allocation strategies, with in particular support for allocating temporary objects using Bumper.jl. The interface for working with backends was introduced in v4, and it is described as allowing dynamic switching between different implementations of the primitive tensor operations, or between different strategies for allocating new tensor objects.

That last point is the interesting one. A tensor contraction allocates a fair amount of temporary memory, and how that memory is obtained is not a detail you want hard-coded into a library. Bumper.jl is a GPU memory management package, and being able to route temporary allocations through it means the same contraction code can run on a plain array, on a CuArray with the default allocator, or on a CuArray whose scratch space is drawn from a pool you manage.

The fourth item is the warning that TensorOperations 4.0 contains several breaking changes and cannot generally be expected to be compatible with previous versions. Breaking changes to part of the interface to make it more sustainable for future improvements and extensions is listed as a deliberate goal of v5, not an accident.

The v4 changes explain why the current interface looks the way it does

The v4 notes are worth reading even if you are only starting on the current version, because they explain the shape of the API you inherit.

The @tensor macro now accepts keyword arguments, to facilitate a variety of options that help with debugging, contraction cost and backend selection. That is a small change with a large effect. Choosing which backend to dispatch to is a decision that in v4 became a call-site decision rather than a global one, and being able to ask for contraction cost information at the call site is what makes the choice of backend an informed one rather than a guess.

The interface for custom types has been changed and thoroughly documented, making it easier to know what to implement, and the consequence is that more general element types of tensors are now possible. If you work with tensor elements that are not plain numbers, this is the section that matters, because it tells you what a custom type has to provide.

Two more changes are about removing sharp edges. Experimental support for automatic differentiation was added by adding reverse-mode chainrules, which later became the improved reverse-mode support in v5. And the cache for temporaries was removed, because of its inconsistent and hard-to-predict behaviour under multithreading. The documentation is candid that the new feature of specifying custom allocation strategies can be used to experiment with novel cache-like behaviour in the future.

That last trade is instructive. A temporary cache would help when you run the same contraction repeatedly, but it interacts badly with multithreading, and the project chose correctness and simplicity over the performance win, then left the door open for users who want the cache back with a properly managed allocator.

The GPU dependency became optional rather than required

There is one change in the v4 list that matters most for anyone who is not doing GPU work, and it is easy to miss among the interface changes.

The support for CuArray objects is moved to a package extension, to avoid unnecessary CUDA dependencies for Julia versions >= 1.9. In Julia, a package extension is a mechanism for loading optional code only when the user has the corresponding package installed, so a CPU-only user does not pull in CUDA.jl and does not need a GPU-capable Julia build.

For a scientific package this is a structural decision, not a convenience. The cost of a heavy optional dependency is not only the download, it is the fact that loading it can change the behaviour of unrelated code, force recompilation, and make the package unusable in environments where the GPU stack is absent. Splitting CuArray support into an extension means the base install stays small.

The v5 note about cuTENSOR v2 and more recent CUDA.jl versions builds on that structure. The GPU path is not a fork of the package, it is the same package with an extension loaded.

The repository layout reflects the structure. There is an ext/ directory at the root, which is where Julia package extensions live, alongside src/ for the implementation, test/ for the test suite, and docs/ for the documentation. There is also a .buildkite/ directory, which is the continuous integration service for Julia packages, and the Project.toml that declares the package's own dependencies and compat bounds.

A package that publishes to a registry and a DOI

The distribution story is the standard Julia one, and it comes with a few Julia-specific pieces worth knowing.

The readme is the repository readme for a package that lives in the General registry, so installation is a Pkg operation rather than a build from source. The project homepage points at the hosted documentation, with a stable URL and a separate latest URL for the development docs, which is the standard split so that a reader following a paper link lands on a versioned manual rather than a moving target.

The badge row shows the quality signals the Julia ecosystem expects. There is a CI badge for the continuous integration runs, a PkgEval badge from Nanosoldier, which is the service that runs package evaluation tests across different Julia versions and dependency configurations, a Codecov badge for test coverage, and an Aqua badge from JuliaTesting. Aqua is a package of static quality checks for Julia packages, covering things like method ambiguities, unbound type parameters, stale dependencies and piracy, and its presence in the badge row means the maintainers run those checks on their own code.

There is also a Zenodo DOI badge, 10.5281/zenodo.3245496, which archives each release with a citable identifier. For a package used in published work, that DOI is what makes a specific version of your results reproducible.

The citation guidance points at CITATION.cff in the repository root, or the Cite this repository button on GitHub, which exports BibTeX or APA. The file is in the root listing, so the citation metadata is maintained alongside the code rather than in a separate document.

Where the package is the wrong tool, and where the ecosystem says to go instead

There are two honest limits to what this package does.

The first is scope. TensorOperations.jl is a contraction and index-notation layer, not an array library. If you want a general-purpose n-dimensional array with the full set of linear algebra operations and a large ecosystem of solvers around it, you want a different package, and the Julia ecosystem has several mature ones. TensorOperations sits above whichever array type you use, and the v4 work on custom types and more general element types is what makes that layering work.

The second is the versioning. The documentation states plainly that version 4.0 contains several breaking changes and cannot generally be expected to be compatible with previous versions, and v5 lists breaking changes to part of the interface as a deliberate goal. If you are starting fresh on the 5.x line that is fine. If you have an existing codebase on an older major version, upgrading is a migration with API changes, not a dependency bump, and the readme gives you no deprecation path between the versions.

On the alternative question, the more interesting comparison is not with another Julia package but with writing the contraction by hand. A hand-written nested loop over a 6-index tensor with three contracted indices is exactly the code TensorOperations exists to remove, and the keyword arguments for backend selection and contraction cost exist because the choice of implementation is not always obvious from the expression alone. For a small number of fixed contractions that you control, hand-writing can be faster to get right. For anything you intend to run on a GPU, or to differentiate through, or to change shape, the macro is the better default.

Editorial conclusion

TensorOperations.jl suits a Julia researcher or package author who writes tensor contractions by hand and wants Einstein notation with GPU acceleration and differentiation support without hand-rolling the index bookkeeping. It is the wrong choice if you want a drop-in replacement for a mature array library, since the documentation warns that version 4.0 and the 5.x line contain breaking changes that cannot generally be expected to be compatible with previous versions. Before upgrading across a major version, read the What is new sections in the readme, because the interface changed deliberately, and confirm your backend choice works with your Julia version, since the CuArray support was moved into a package extension to avoid unnecessary CUDA dependencies.

Frequently asked questions

Official sources

1. Issues
2. Project website
3. QuantumKitHub/TensorOperations.jl on GitHub
4. README
5. Releases