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