extended-einsum 0.1.0__py3-none-any.whl

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/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/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/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]