backend.pyi

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, ClassVar, Iterable, Mapping, Protocol

__all__ = [
    "MathematicsBackend",
    "BackendUnavailableError",
    "register_backend",
    "get_backend",
    "available_backends",
    "ensure_array",
    "ensure_numpy",
]

class BackendUnavailableError(RuntimeError): ...

class MathematicsBackend(Protocol):
    name: str
    supports_autodiff: bool
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...

BackendFactory = Callable[[], MathematicsBackend]

@dataclass
class _NumpyBackend:
    name: ClassVar[str] = ...
    supports_autodiff: ClassVar[bool] = ...
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...

@dataclass
class _JaxBackend:
    name: ClassVar[str] = ...
    supports_autodiff: ClassVar[bool] = ...
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...

@dataclass
class _TorchBackend:
    name: ClassVar[str] = ...
    supports_autodiff: ClassVar[bool] = ...
    def as_array(self, value: Any, *, dtype: Any | None = None) -> Any: ...
    def eig(self, matrix: Any) -> tuple[Any, Any]: ...
    def eigh(self, matrix: Any) -> tuple[Any, Any]: ...
    def matrix_exp(self, matrix: Any) -> Any: ...
    def norm(
        self, value: Any, *, ord: Any | None = None, axis: Any | None = None
    ) -> Any: ...
    def einsum(self, pattern: str, *operands: Any, **kwargs: Any) -> Any: ...
    def matmul(self, a: Any, b: Any) -> Any: ...
    def conjugate_transpose(self, matrix: Any) -> Any: ...
    def stack(self, arrays: Iterable[Any], *, axis: int = 0) -> Any: ...
    def to_numpy(self, value: Any) -> Any: ...

def ensure_array(
    value: Any, *, dtype: Any | None = None, backend: MathematicsBackend | None = None
) -> Any: ...
def ensure_numpy(value: Any, *, backend: MathematicsBackend | None = None) -> Any: ...
def register_backend(
    name: str,
    factory: BackendFactory,
    *,
    aliases: Iterable[str] | None = None,
    override: bool = False,
) -> None: ...
def get_backend(name: str | None = None) -> MathematicsBackend: ...
def available_backends() -> Mapping[str, BackendFactory]: ...