extended-einsum 0.1.0__tar.gz

extended_einsum-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 FSU Theoretical Computer Science II
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

extended_einsum-0.1.0/PKG-INFO

@@ -0,0 +1,139 @@
+Metadata-Version: 2.4
+Name: extended-einsum
+Version: 0.1.0
+Summary: Build, optimize, and execute numerically stable extended Einstein summation expressions with PyTorch, NumPy, and optional JAX.
+Keywords: einsum,machine-learning,probabilistic circuits,tensor
+Author: Christoph Staudt, Maurice Wenig
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Dist: numpy>=2.1.6
+Requires-Dist: sesum>=0.4.0
+Requires-Dist: torch>=2.12.0
+Requires-Dist: jax>=0.10.1 ; extra == 'jax'
+Requires-Dist: jax[cuda13]>=0.10.1 ; extra == 'jax-cuda'
+Requires-Dist: matplotlib>=3.10.9 ; extra == 'visualization'
+Requires-Dist: networkx>=3.6.1 ; extra == 'visualization'
+Requires-Python: >=3.12
+Project-URL: Homepage, https://github.com/ti2-group/extended_einsum
+Project-URL: Repository, https://github.com/ti2-group/extended_einsum
+Project-URL: Issues, https://github.com/ti2-group/extended_einsum/issues
+Project-URL: Changelog, https://github.com/ti2-group/extended_einsum/blob/main/CHANGELOG.md
+Provides-Extra: jax
+Provides-Extra: jax-cuda
+Provides-Extra: visualization
+Description-Content-Type: text/markdown
+
+# extended-einsum
+
+[![CI](https://github.com/ti2-group/extended_einsum/actions/workflows/ci.yml/badge.svg)](https://github.com/ti2-group/extended_einsum/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/extended-einsum.svg)](https://pypi.org/project/extended-einsum/)
+[![Python](https://img.shields.io/pypi/pyversions/extended-einsum.svg)](https://pypi.org/project/extended-einsum/)
+[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/ti2-group/extended_einsum/blob/main/LICENSE)
+
+`extended-einsum` builds tensor-expression graphs, rewrites contraction programs, and executes them with numerically stable evaluation strategies. PyTorch is the primary backend; NumPy is included in the base installation and JAX is optional.
+
+The package is an expression engine, not a drop-in replacement for `torch.einsum` or `numpy.einsum`. Wrap backend arrays with `xe.array`, compose operations lazily, then call `materialize`.
+
+> **Status:** `0.1.0` is an alpha release. The public expression API is usable, but indexing and some operators or stability-mode combinations are not implemented yet.
+
+## Installation
+
+Python 3.12, 3.13, and 3.14 are supported.
+
+```bash
+pip install extended-einsum
+```
+
+Optional features are installed explicitly:
+
+```bash
+pip install "extended-einsum[jax]"
+pip install "extended-einsum[jax-cuda]"
+pip install "extended-einsum[visualization]"
+```
+
+The CUDA extra follows JAX's CUDA 13 installation. PyTorch device support is determined by the PyTorch build available from your configured package index.
+
+## Quickstart
+
+```python
+import torch
+import extended_einsum as xe
+
+left_torch = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+right_torch = torch.tensor([[2.0, 0.0], [1.0, 2.0]])
+
+left = xe.array(left_torch)
+right = xe.array(right_torch)
+expression = xe.einsum("ik,kj->ij", xe.exp(left), right)
+result = expression.materialize(stability_mode="unstable")
+
+assert torch.allclose(
+    result.backend_array,
+    torch.einsum("ik,kj->ij", torch.exp(left_torch), right_torch),
+)
+```
+
+The same interface accepts NumPy arrays. With the `jax` extra installed it also accepts JAX arrays:
+
+```python
+import numpy as np
+import extended_einsum as xe
+
+source = xe.array(np.arange(6.0).reshape(2, 3))
+result = xe.softmax(source, axis=1).materialize()
+print(result.backend_array)
+```
+
+## Expressions and stability
+
+The public interface currently includes:
+
+- `array`, `einsum`, `stack`, `take`, `slice`, and `select`
+- `exp`, `log`, `sin`, `cos`, `tan`, `sqrt`, `inverse`, and `softmax`
+- `TensorExpression.materialize` and `extract_program`
+
+Supported execution backends are PyTorch, NumPy, and optional JAX. Rich programs can be evaluated in `unstable`, `scaled_min`, `scaled_sum`, `logspace_min`, or `logspace_max` mode. Support is operator-dependent; unsupported combinations raise `NotImplementedError` instead of silently changing semantics.
+
+The preprocessing API also provides expression folding and contraction-path optimization. DAG plotting is available from `extended_einsum.visualization` when the visualization extra is installed.
+
+## Examples
+
+All introductory scripts in [`examples/`](https://github.com/ti2-group/extended_einsum/tree/main/examples) are standalone: copy one into another project with `extended-einsum` installed and run it from any working directory.
+
+```bash
+python examples/torch_quickstart.py
+python examples/numpy_backend.py
+python examples/stable_materialization.py
+python examples/jax_backend.py                    # requires [jax]
+python examples/visualize_expression.py graph.png # requires [visualization]
+```
+
+The advanced [`demo/cirkit.py`](https://github.com/ti2-group/extended_einsum/blob/main/demo/cirkit.py) integration is also standalone, but requires the repository demo dependencies and currently Python 3.12 or 3.13 because of Cirkit's SciPy constraint. Its defaults use a small synthetic circuit; training and benchmark sweeps are opt-in. Generated benchmark data and plots are repository artifacts and are never included in distributions.
+
+## Development
+
+This project uses [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv sync --group dev --group demo --extra jax --extra visualization
+uv run ruff check .
+uv run pytest
+uv build
+```
+
+See [CONTRIBUTING.md](https://github.com/ti2-group/extended_einsum/blob/main/CONTRIBUTING.md) for development details and [PUBLISHING.md](https://github.com/ti2-group/extended_einsum/blob/main/PUBLISHING.md) for the release process.
+
+## License
+
+Copyright © 2026 FSU Theoretical Computer Science II. Distributed under the [MIT License](https://github.com/ti2-group/extended_einsum/blob/main/LICENSE).

extended_einsum-0.1.0/README.md

@@ -0,0 +1,104 @@
+# extended-einsum
+
+[![CI](https://github.com/ti2-group/extended_einsum/actions/workflows/ci.yml/badge.svg)](https://github.com/ti2-group/extended_einsum/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/extended-einsum.svg)](https://pypi.org/project/extended-einsum/)
+[![Python](https://img.shields.io/pypi/pyversions/extended-einsum.svg)](https://pypi.org/project/extended-einsum/)
+[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/ti2-group/extended_einsum/blob/main/LICENSE)
+
+`extended-einsum` builds tensor-expression graphs, rewrites contraction programs, and executes them with numerically stable evaluation strategies. PyTorch is the primary backend; NumPy is included in the base installation and JAX is optional.
+
+The package is an expression engine, not a drop-in replacement for `torch.einsum` or `numpy.einsum`. Wrap backend arrays with `xe.array`, compose operations lazily, then call `materialize`.
+
+> **Status:** `0.1.0` is an alpha release. The public expression API is usable, but indexing and some operators or stability-mode combinations are not implemented yet.
+
+## Installation
+
+Python 3.12, 3.13, and 3.14 are supported.
+
+```bash
+pip install extended-einsum
+```
+
+Optional features are installed explicitly:
+
+```bash
+pip install "extended-einsum[jax]"
+pip install "extended-einsum[jax-cuda]"
+pip install "extended-einsum[visualization]"
+```
+
+The CUDA extra follows JAX's CUDA 13 installation. PyTorch device support is determined by the PyTorch build available from your configured package index.
+
+## Quickstart
+
+```python
+import torch
+import extended_einsum as xe
+
+left_torch = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+right_torch = torch.tensor([[2.0, 0.0], [1.0, 2.0]])
+
+left = xe.array(left_torch)
+right = xe.array(right_torch)
+expression = xe.einsum("ik,kj->ij", xe.exp(left), right)
+result = expression.materialize(stability_mode="unstable")
+
+assert torch.allclose(
+    result.backend_array,
+    torch.einsum("ik,kj->ij", torch.exp(left_torch), right_torch),
+)
+```
+
+The same interface accepts NumPy arrays. With the `jax` extra installed it also accepts JAX arrays:
+
+```python
+import numpy as np
+import extended_einsum as xe
+
+source = xe.array(np.arange(6.0).reshape(2, 3))
+result = xe.softmax(source, axis=1).materialize()
+print(result.backend_array)
+```
+
+## Expressions and stability
+
+The public interface currently includes:
+
+- `array`, `einsum`, `stack`, `take`, `slice`, and `select`
+- `exp`, `log`, `sin`, `cos`, `tan`, `sqrt`, `inverse`, and `softmax`
+- `TensorExpression.materialize` and `extract_program`
+
+Supported execution backends are PyTorch, NumPy, and optional JAX. Rich programs can be evaluated in `unstable`, `scaled_min`, `scaled_sum`, `logspace_min`, or `logspace_max` mode. Support is operator-dependent; unsupported combinations raise `NotImplementedError` instead of silently changing semantics.
+
+The preprocessing API also provides expression folding and contraction-path optimization. DAG plotting is available from `extended_einsum.visualization` when the visualization extra is installed.
+
+## Examples
+
+All introductory scripts in [`examples/`](https://github.com/ti2-group/extended_einsum/tree/main/examples) are standalone: copy one into another project with `extended-einsum` installed and run it from any working directory.
+
+```bash
+python examples/torch_quickstart.py
+python examples/numpy_backend.py
+python examples/stable_materialization.py
+python examples/jax_backend.py                    # requires [jax]
+python examples/visualize_expression.py graph.png # requires [visualization]
+```
+
+The advanced [`demo/cirkit.py`](https://github.com/ti2-group/extended_einsum/blob/main/demo/cirkit.py) integration is also standalone, but requires the repository demo dependencies and currently Python 3.12 or 3.13 because of Cirkit's SciPy constraint. Its defaults use a small synthetic circuit; training and benchmark sweeps are opt-in. Generated benchmark data and plots are repository artifacts and are never included in distributions.
+
+## Development
+
+This project uses [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv sync --group dev --group demo --extra jax --extra visualization
+uv run ruff check .
+uv run pytest
+uv build
+```
+
+See [CONTRIBUTING.md](https://github.com/ti2-group/extended_einsum/blob/main/CONTRIBUTING.md) for development details and [PUBLISHING.md](https://github.com/ti2-group/extended_einsum/blob/main/PUBLISHING.md) for the release process.
+
+## License
+
+Copyright © 2026 FSU Theoretical Computer Science II. Distributed under the [MIT License](https://github.com/ti2-group/extended_einsum/blob/main/LICENSE).

extended_einsum-0.1.0/pyproject.toml

@@ -0,0 +1,85 @@
+[project]
+name = "extended-einsum"
+version = "0.1.0"
+description = "Build, optimize, and execute numerically stable extended Einstein summation expressions with PyTorch, NumPy, and optional JAX."
+readme = "README.md"
+requires-python = ">=3.12"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [
+    { name = "Christoph Staudt" },
+    { name = "Maurice Wenig" },
+]
+keywords = [
+    "einsum",
+    "machine-learning",
+    "probabilistic circuits",
+    "tensor",
+]
+classifiers = [
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Typing :: Typed",
+]
+dependencies = [
+    "numpy>=2.1.6",
+    "sesum>=0.4.0",
+    "torch>=2.12.0",
+]
+
+[project.optional-dependencies]
+jax = ["jax>=0.10.1"]
+jax-cuda = ["jax[cuda13]>=0.10.1"]
+visualization = [
+    "matplotlib>=3.10.9",
+    "networkx>=3.6.1",
+]
+
+[project.urls]
+Homepage = "https://github.com/ti2-group/extended_einsum"
+Repository = "https://github.com/ti2-group/extended_einsum"
+Issues = "https://github.com/ti2-group/extended_einsum/issues"
+Changelog = "https://github.com/ti2-group/extended_einsum/blob/main/CHANGELOG.md"
+
+[build-system]
+requires = ["uv_build>=0.11.19,<0.12.0"]
+build-backend = "uv_build"
+
+[tool.uv]
+constraint-dependencies = [
+    "numpy>=2.3.0; python_version >= '3.14'",
+    "scipy>=1.17.0; python_version >= '3.14'",
+]
+
+[dependency-groups]
+demo = [
+    "libcirkit>=0.2.1; python_version < '3.14'",
+    "matplotlib>=3.10.9",
+    "networkx>=3.6.1",
+    "pandas>=2.3.3",
+    "seaborn>=0.13.2",
+    "torchvision>=0.27.0",
+]
+dev = [
+    "pytest>=9.1.1",
+    "ruff>=0.6.0",
+]
+
+[tool.ruff]
+line-length = 200
+target-version = "py312"
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+line-ending = "lf"

extended_einsum-0.1.0/src/extended_einsum/__init__.py

@@ -0,0 +1,43 @@
+from importlib.metadata import PackageNotFoundError, version
+
+from extended_einsum.interface import TensorExpression as TensorExpression
+from extended_einsum.interface import array as array
+from extended_einsum.interface import cos as cos
+from extended_einsum.interface import einsum as einsum
+from extended_einsum.interface import exp as exp
+from extended_einsum.interface import extract_program as extract_program
+from extended_einsum.interface import inverse as inverse
+from extended_einsum.interface import log as log
+from extended_einsum.interface import select as select
+from extended_einsum.interface import sin as sin
+from extended_einsum.interface import slice as slice
+from extended_einsum.interface import softmax as softmax
+from extended_einsum.interface import sqrt as sqrt
+from extended_einsum.interface import stack as stack
+from extended_einsum.interface import take as take
+from extended_einsum.interface import tan as tan
+
+try:
+    __version__ = version("extended-einsum")
+except PackageNotFoundError:
+    __version__ = "0.1.0.dev0"
+
+__all__ = [
+    "TensorExpression",
+    "__version__",
+    "array",
+    "cos",
+    "einsum",
+    "exp",
+    "extract_program",
+    "inverse",
+    "log",
+    "select",
+    "sin",
+    "slice",
+    "softmax",
+    "sqrt",
+    "stack",
+    "take",
+    "tan",
+]

extended_einsum-0.1.0/src/extended_einsum/backend_translation/__init__.py

@@ -0,0 +1,7 @@
+from .backend import BackendArray as BackendArray
+from .backend import BackendCompiler as BackendCompiler
+from .backend import BackendFunctions as BackendFunctions
+from .backend import BackendProgram as BackendProgram
+from .backend import get_backend_of_array as get_backend_of_array
+from .runtime import run_program as run_program
+from .translate import translate_to_backend_program as translate_to_backend_program

extended_einsum-0.1.0/src/extended_einsum/backend_translation/backend.py

@@ -0,0 +1,110 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Callable, Generic, Protocol, TypeVar
+
+import numpy as np
+import torch
+
+from extended_einsum.language.types import Backend, HasShape
+
+
+class BackendArray(Protocol):
+    @property
+    def shape(self) -> tuple[int, ...] | torch.Size: ...
+
+
+TBackendArray = TypeVar("TBackendArray", bound=BackendArray)
+
+
+class BackendFunctions(Protocol[TBackendArray]):
+    @staticmethod
+    def stop_gradient(array: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def exp(array: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def log(array: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def sum(array: TBackendArray, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> TBackendArray: ...
+
+    @staticmethod
+    def max(array: TBackendArray, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> TBackendArray: ...
+
+    @staticmethod
+    def min(array: TBackendArray, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> TBackendArray: ...
+
+    @staticmethod
+    def maximum(array_1: TBackendArray, array_2: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def reshape(array: TBackendArray, shape: tuple[int, ...]) -> TBackendArray: ...
+
+    @staticmethod
+    def broadcast_to(array: TBackendArray, shape: tuple[int, ...]) -> TBackendArray: ...
+
+    @staticmethod
+    def stack(arrays: Sequence[TBackendArray], axis: int) -> TBackendArray: ...
+
+    @staticmethod
+    def concat(arrays: Sequence[TBackendArray], axis: int) -> TBackendArray: ...
+
+    @staticmethod
+    def take(array: TBackendArray, indices: TBackendArray, axis: int) -> TBackendArray: ...
+
+    @staticmethod
+    def select(array: TBackendArray, axis: int, index: int) -> TBackendArray: ...
+
+    @staticmethod
+    def slice(array: TBackendArray, start: int, stop: int, axis: int) -> TBackendArray: ...
+
+    @staticmethod
+    def softmax(array: TBackendArray, axis: int | tuple[int, ...]) -> TBackendArray: ...
+
+    @staticmethod
+    def einsum(format_string: str, *operands: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def add(summand_array_1: TBackendArray, summand_array_2: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def subtract(minuend_array: TBackendArray, subtrahend_array: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def multiply(factor_array_1: TBackendArray, factor_array_2: TBackendArray) -> TBackendArray: ...
+
+    @staticmethod
+    def divide(dividend_array: TBackendArray, divisor_array: TBackendArray) -> TBackendArray: ...
+
+
+@dataclass(frozen=True)
+class BackendProgram(Generic[TBackendArray]):
+    backend_calls: list[Callable[[Sequence[TBackendArray]], TBackendArray]]
+    call_arguments: list[tuple[int, ...]]
+    n_inputs: int
+
+
+class BackendCompiler(Protocol[TBackendArray]):
+    @staticmethod
+    def compile(
+        program: BackendProgram[TBackendArray],
+        inputs: Sequence[TBackendArray],
+    ) -> Callable[[Sequence[TBackendArray]], TBackendArray]: ...
+
+
+def get_backend_of_array(array: HasShape) -> Backend:
+    if isinstance(array, torch.Tensor):
+        return "torch"
+    elif isinstance(array, np.ndarray):
+        return "numpy"
+
+    try:
+        import jax
+    except ModuleNotFoundError:
+        jax = None  # type: ignore[assignment]
+
+    if jax is not None and isinstance(array, jax.Array):
+        return "jax"
+
+    raise ValueError(f"Unsupported array type: {type(array)}")

extended_einsum-0.1.0/src/extended_einsum/backend_translation/runtime.py

@@ -0,0 +1,21 @@
+from collections.abc import Sequence
+from typing import TypeVar
+
+from extended_einsum.backend_translation.backend import BackendArray, BackendProgram
+
+TBackendArray = TypeVar("TBackendArray", bound=BackendArray)
+
+
+def run_program(
+    program: BackendProgram[TBackendArray],
+    inputs: Sequence[TBackendArray],
+) -> TBackendArray:
+    if len(inputs) != program.n_inputs:
+        raise ValueError(f"The number of inputs ({len(inputs)}) does not match the number of inputs ({program.n_inputs}) in the program.")
+
+    tensors: list[TBackendArray] = list(inputs)
+    for backend_call, argument_ids in zip(program.backend_calls, program.call_arguments):
+        argument_tensors = [tensors[argument] for argument in argument_ids]
+        result = backend_call(argument_tensors)
+        tensors.append(result)
+    return tensors[-1]