OMEinsum

You can specify tensor operations using either the EinCode object or the @ein_str string literal. Both methods are equivalent and allow you to define the index mapping for contractions.

Use EinCode when you want to programmatically construct the index mapping. Use the string literal (via the ein"..." syntax) for a more concise, readable notation.

using OMEinsum

# Using the string literal
code1 = ein"ij,jk -> ik"

# Using the EinCode object
ixs = [[1, 2], [2, 3]]
iy = [1, 3]
code2 = EinCode(ixs, iy)

OMEinsum provides two distinct backends for GPU tensor contractions via CUDA:

DefaultBackend works by reshaping tensors into matrix multiplication (GEMM) formats, while CuTensorBackend provides native tensor contraction without reshape/permute overhead.

Since v0.4, OMEinsum no longer optimizes the contraction order automatically. To ensure efficiency or to help the engine select specific kernels, you must specify the order manually using brackets in the string or macro syntax.

String style: ein"is, (oij, js) -> os"(x, W, y)

Macro style: @ein Z[o,s] := x[i,s] * (W[o,i,j] * y[j,s]);

If you do not provide brackets, the package may fall back to loop_einsum, which is slower. You can use allow_loops(false) to force an error if a fallback is used, helping you identify where manual ordering is needed.

# Manual ordering via brackets
Z = ein"is, (oij, js) -> os"(x, W, y);

# Enforce manual ordering by disabling loop fallback
allow_loops(false);
Zl = ein"is, oij, js -> os"(x, W, y); # This will throw an error

To install OMEinsum in Julia (version 1.5 or higher), use the package manager in the REPL:

1. Enter the package REPL by typing ].
2. Run add OMEinsum.

pkg> add OMEinsum

OMEinsum supports CUDA GPU acceleration via two distinct backends. Choosing the right one depends on your contraction pattern:

1. DefaultBackend() (CUBLAS): Best for standard matrix-like (GEMM) contractions.
2. CuTensorBackend() (cuTENSOR): Best for general tensor contractions and non-GEMM patterns (e.g., tensor networks). It provides native contraction without reshape/permute overhead.

To switch backends, use set_einsum_backend!().

Note: Using CuTensorBackend requires the cuTENSOR package (which provides the CuTENSORExt extension).

using CUDA
using cuTENSOR
using OMEinsum

# Switch to cuTENSOR for general tensor contractions
set_einsum_backend!(CuTensorBackend())
C = ein"ijk,jkl->il"(A, B)

When performing contractions with more than two input tensors, the standard @ein_str macro does not optimize the contraction order. To ensure efficient computation for large N-ary contractions, use one of the following:

1. optein"..." string literal: Automatically optimizes the contraction order without requiring explicit tensor sizes.
2. Manual specification: Use ein"(indices),(indices)->output" to manually specify the contraction grouping (e.g., grouping specific tensors together).

using OMEinsum
tensors = [randn(100, 100) for _ in 1:4]

# Optimized contraction (automatic)
optein"ij,jk,kl,lm->im"(tensors...)

# Manually specified contraction order
ein"(ij,jk),(kl,lm)->im"(tensors...)

When using the @ein_str literal for more than two input tensors, the contraction order is not automatically optimized. To optimize the order, you can construct an EinCode object and pass it to optimize_code.

This function requires a size_dict (mapping labels to their sizes) and an optimization algorithm. A common choice is TreeSA, a local search algorithm. The result is a binary contraction tree of type SlicedEinsum or NestedEinsum.

using OMEinsum, OMEinsumContractionOrders, OMEinsumContractionOrders.Graphs

# 1. Construct an EinCode from a graph
graph = random_regular_graph(20, 3; seed=42)
code = EinCode([[e.src, e.dst] for e in edges(graph)], Int[])

# 2. Define tensor sizes
size_dict = uniformsize(code, 2)

# 3. Optimize the contraction order
optcode = optimize_code(code, size_dict, TreeSA(ntrials=1))

To use OMEinsum.jl, import the package and use the ein macro to define your tensor contraction operation. You can then use optimize_code to find an efficient contraction order using a specific optimizer (like TreeSA()) and a size hint (via uniformsize). The resulting optimized code object can be called like a function with the input tensors as arguments.

using OMEinsum
code = ein"ij,jk,kl,lm->im" # define the einsum operation
optcode = optimize_code(code, uniformsize(code, 100), TreeSA())  # optimize the contraction order
optcode(randn(100, 100), randn(100, 100), randn(100, 100), randn(100, 100))  # compute the result

You can accelerate tensor contractions by uploading your data to the GPU using CUDA.jl. Once data is on the GPU (e.g., as a CuArray), OMEinsum will perform the contraction on the device. For optimal performance, use CUDA.@sync when benchmarking or executing operations to ensure synchronization.

using CUDA, OMEinsum

# Prepare data
A, B = rand(1000, 1000), rand(1000, 1000)
code = ein"ij,jk->ik"

# Upload to GPU
cuA, cuB = CuArray(A), CuArray(B)

# Execute on GPU
result = CUDA.@sync code(cuA, cuB)

If the memory usage (space complexity) of a contraction is too high, you can use slicing to reduce it. Slicing is performed using the slice_code function and a slicer object (e.g., TreeSASlicer).

Slicing involves a trade-off: it reduces space complexity at the cost of increasing time complexity.

# Assuming optcode and size_dict are already defined
slicer = TreeSASlicer(score=ScoreFunction(sc_target=2))
scode = slice_code(optcode, size_dict, slicer)

# scode is a SlicedEinsum object