sesum 0.4.1

Scalable einsum with optimized path computation and efficient execution of large, complex tensor expressions

Sesum

Sesum (Scalable einsum) combines a fast C/C++ backend (with AVX2, AVX-512, and ARM64 support) and a clean Python interface to efficiently compute contraction paths and execute einsum expressions.

Sesum combines features rarely found in a single einsum library:

• Single library: Efficient path computation and execution in one package, requiring only NumPy.
• Highly scalable: Handles large-scale expressions with thousands of tensors.
• Large integer support: Supports int128, uint128, and unlimited-precision bigint.
• Sparse backend: Enables sparse execution for general tensor shapes, with an internal sparse representation optimized for sparse contractions.
• Progress and debug info: Optional progress bars and detailed debug output.
• Flexible indexing: Supports repeating indices in both input and output (e.g. "ijji,kki,lkll->iilki").
• Semiring support: Allows replacing standard + and × with max and + (max-plus semiring) for optimization-style tasks.

Installation

Sesum is available on PyPI:

pip install sesum

Usage Example

The example below demonstrates basic usage of Sesum.

import numpy as np
import sesum as sr

# define an einsum expression
format_string = "dgfh,bgdec,ec,chfa->ab"
arrays = [np.random.rand(*shape) for shape in [(2, 2, 2, 2), (2, 2, 2, 2, 2), (2, 2), (2, 2, 2, 2)]]

# compute the contraction path for the einsum expression
path, flops_log10, size_log2 = sr.compute_path(format_string, *arrays)

# execute the einsum expression
result = sr.sesum(format_string, *arrays, path=path)

Use help(sr.compute_path) to view all parameters and their descriptions for path computation, and help(sr.sesum) for einsum execution.

It is recommended to review all supported parameters at least once to fully understand and utilize the capabilities of Sesum. For example, a more realistic usage for finding efficient contraction paths might look as follows:

path, flops_log10, size_log2 = sr.compute_path(
    format_string, *arrays,
    seed=0,                      # random seed for reproducibility
    minimize="flops",            # objective to minimize: "flops", "size" or "b-flops"
    algorithm="kahypar",         # path algorithm: "greedy" or "kahypar"
    skops_alpha=0,               # >0 (e.g. 64) may speed up execution, but adds flops
    max_repeats=128,             # max number of optimization repeats
    max_time=0.0,                # time limit in seconds (0.0 = unlimited)
    progbar=True,                # show progress bar during optimization
    is_outer_optimal=False,      # consider outer products in optimal search
    threshold_optimal=18,        # max number of inputs for optimal search
    threads=0,                   # number of threads (0 = use all cores)
    is_linear=True               # return linear (vs SSA) contraction path
)

As with path computation, there are numerous options available to configure einsum execution. The example below demonstrates how to enforce the use of np.float32 as the computation data type.

result = sr.sesum(
    format_string, *arrays,
    path=path,                  # precomputed contraction path
    dtype=np.float32,           # force np.float32, supports custom types like sr.bigint
    debug=True,                 # print debug info during execution
    safe_convert=False,         # if True, avoid overflow in input type conversion
    backend="dense",            # "dense" for general use, "sparse" for sparse-aware execution on general tensor shapes
    semiring=sr.standard,       # semiring used: sr.standard or sr.max_plus
    output="dense",             # "dense" returns np.ndarray, "sparse" returns SesumSparseTensor for sr.standard
    postorder=False             # False follows path order, True uses postorder tree traversal
)

By default Sesum executes pairwise contractions in the provided or generated path order. Set postorder=True only when you explicitly want postorder tree traversal.

Sparse Tensors and Sparse Backend

Sesum provides a native sparse tensor type for the sparse backend. The public sparse constructor is sparse_coo_tensor, which follows the PyTorch-style COO layout with indices shaped as (ndim, nnz):

S = sr.sparse_coo_tensor(
    [[0, 1, 1], [2, 0, 0]],
    [5, 7, 3],
    size=(2, 3),
    dtype=np.int64,
)

Coordinates do not need to be sorted. Duplicate coordinates are supported and are coalesced by summing their values, so the example above stores two entries:

S.nnz       # 2
S.coords() # array([[0, 2], [1, 0]])
S.values() # array([5, 10])

coords() returns coordinates in Sesum's sorted/coalesced storage order, which may differ from the input order.

Sesum also accepts (nnz, ndim) coordinate tuples when the shape is unambiguous:

S = sr.sparse_coo_tensor(
    [(0, 2), (1, 0), (1, 0)],
    [5, 7, 3],
    size=(2, 3),
    dtype=np.int64,
)

To convert a dense NumPy array to Sesum's sparse format, use dense_to_sparse. This scans the dense input, so sparse_coo_tensor is usually better when coordinates are already known:

a = np.array([[0, 5, 0],
              [7, 0, 0]], dtype=np.float64)

S = sr.dense_to_sparse(a)

If duplicate values sum to zero, the entry is removed. For sr.int128, sr.uint128, and sr.bigint, values() returns a NumPy object array containing Python int values.

Sparse tensors can be used directly in contractions:

B = np.ones((3, 2), dtype=np.int64)

dense_result = sr.sesum("ij,jk->ik", S, B, backend="sparse")
sparse_result = sr.sesum("ij,jk->ik", S, B, backend="sparse", output="sparse")

type(sparse_result)  # sr.SesumSparseTensor
sparse_result.to_numpy()

output="sparse" may also be used with backend="dense" for the standard semiring. In that case Sesum computes the contraction densely and converts the dense result to a native sparse tensor before returning it.

Dense and sparse inputs may be mixed. Dtype inference follows dense Sesum: if sparse input dtypes differ from the inferred or requested execution dtype, Sesum converts them sparsely using SesumSparseTensor.astype(dtype) rather than densifying the tensor.

S32 = sr.dense_to_sparse(np.array([[1, 0]], dtype=np.int32))
T64 = sr.dense_to_sparse(np.array([[2], [3]], dtype=np.int64))

out = sr.sesum("ij,jk->ik", S32, T64, backend="sparse")
out.dtype  # np.int64

The sparse backend currently supports the standard semiring for native sparse inputs and sparse outputs. The dense backend supports both sr.standard and sr.max_plus.

References

The multiple-cost-functions strategy employed in the greedy contraction path algorithm is discussed in the paper, titled "Optimizing Tensor Contraction Paths: A Greedy Algorithm Approach With Improved Cost Functions".

@article{OrglerB24,
  author    = {Sheela Orgler and Mark Blacher},
  title     = {{Optimizing Tensor Contraction Paths: A Greedy Algorithm Approach With Improved Cost Functions}},
  year      = {2024},
  journal   = {{arXiv}}
}

The algorithm driving the hypergraph partitioning approach for finding efficient contraction paths is discussed in the paper titled "Improved cut strategy for tensor network contraction orders".

@inproceedings{StaudtBKLG24,
  author       = {Christoph Staudt and Mark Blacher and Julien Klaus and Farin Lippmann and Joachim Giesen},
  title        = {{Improved Cut Strategy for Tensor Network Contraction Orders}},
  booktitle    = {{SEA}},
  year         = {2024}
}

Sesum is designed to execute large einsum problems like those introduced in "Einsum Benchmark: Enabling the Development of Next-Generation Tensor Execution Engines".

@inproceedings{BlacherSKWMBELG24,
  author       = {Mark Blacher and Christoph Staudt and Julien Klaus and Maurice Wenig and Niklas Merk and Alexander Breuer and Max Engel and S{\"{o}}ren Laue and Joachim Giesen},
  title        = {{Einsum Benchmark: Enabling the Development of Next-Generation Tensor Execution Engines}},
  booktitle    = {{NeurIPS}},
  year         = {2024}
}

The sparse backend was utilized in the experiments of "Exploiting Dynamic Sparsity in Einsum".

@inproceedings{StaudtBHKBG25,
  author       = {Christoph Staudt and Mark Blacher and Tim Hoffmann and Kaspar Kasche and Olaf Beyersdorff and Joachim Giesen},
  title        = {{Exploiting Dynamic Sparsity in Einsum}},
  booktitle    = {{NeurIPS}},
  year         = {2025}
}

Source Code

The source code used to compile Sesum is bundled with the installed Python package and accessible from its installation directory.