Metadata-Version: 2.1
Name: sesum
Version: 0.4.1
Summary: Scalable einsum with optimized path computation and efficient execution of large, complex tensor expressions
Author: Mark Blacher
License: GPL-3.0
Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
Classifier: Programming Language :: Python :: 3
Description-Content-Type: text/markdown
Requires-Dist: numpy

# **Sesum**

**Sesum** (**S**calable **e**in**sum**) 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:

```bash
pip install sesum
```

## Usage Example

The example below demonstrates basic usage of Sesum.


```python
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:

```python
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.

```python
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)`:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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.

```python
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"_](https://arxiv.org/pdf/2405.09644).

```bibtex
@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"_](https://drops.dagstuhl.de/storage/00lipics/lipics-vol301-sea2024/LIPIcs.SEA.2024.27/LIPIcs.SEA.2024.27.pdf).

```bibtex
@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"_](https://openreview.net/pdf?id=tllpLtt14h).

```bibtex
@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"_](https://openreview.net/pdf?id=ixOpURt7wC).

```bibtex
@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.
