1"""Types used in the opt_einsum package."""
2
3from collections import namedtuple
4from typing import Any, Callable, Collection, Dict, FrozenSet, List, Literal, Optional, Tuple, Union
5
6TensorShapeType = Tuple[int, ...]
7PathType = Collection[TensorShapeType]
8
9ArrayType = Any
10
11ArrayIndexType = FrozenSet[str]
12ArrayShaped = namedtuple("ArrayShaped", ["shape"])
13
14ContractionListType = List[Tuple[Any, ArrayIndexType, str, Optional[Tuple[str, ...]], Union[str, bool]]]
15PathSearchFunctionType = Callable[[List[ArrayIndexType], ArrayIndexType, Dict[str, int], Optional[int]], PathType]
16
17# Contract kwargs
18OptimizeKind = Union[
19 None,
20 bool,
21 Literal[
22 "optimal", "dp", "greedy", "random-greedy", "random-greedy-128", "branch-all", "branch-2", "auto", "auto-hq"
23 ],
24 PathType,
25 PathSearchFunctionType,
26]
27BackendType = Literal["auto", "object", "autograd", "cupy", "dask", "jax", "theano", "tensorflow", "torch", "libjax"]