Spaces:
Running
Running
File size: 3,847 Bytes
dc2106c |
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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import sys
from typing import List, TypeVar, Optional, Any, overload, Union, Tuple, Sequence
from numpy import (
ndarray,
dtype,
bool_,
unsignedinteger,
signedinteger,
floating,
complexfloating,
number,
_OrderKACF,
)
from numpy.typing import (
_ArrayLikeBool_co,
_ArrayLikeUInt_co,
_ArrayLikeInt_co,
_ArrayLikeFloat_co,
_ArrayLikeComplex_co,
_DTypeLikeBool,
_DTypeLikeUInt,
_DTypeLikeInt,
_DTypeLikeFloat,
_DTypeLikeComplex,
_DTypeLikeComplex_co,
)
if sys.version_info >= (3, 8):
from typing import Literal
else:
from typing_extensions import Literal
_ArrayType = TypeVar(
"_ArrayType",
bound=ndarray[Any, dtype[Union[bool_, number[Any]]]],
)
_OptimizeKind = Union[
None, bool, Literal["greedy", "optimal"], Sequence[Any]
]
_CastingSafe = Literal["no", "equiv", "safe", "same_kind"]
_CastingUnsafe = Literal["unsafe"]
__all__: List[str]
# TODO: Properly handle the `casting`-based combinatorics
# TODO: We need to evaluate the content `__subscripts` in order
# to identify whether or an array or scalar is returned. At a cursory
# glance this seems like something that can quite easilly be done with
# a mypy plugin.
# Something like `is_scalar = bool(__subscripts.partition("->")[-1])`
@overload
def einsum(
__subscripts: str,
*operands: _ArrayLikeBool_co,
out: None = ...,
dtype: Optional[_DTypeLikeBool] = ...,
order: _OrderKACF = ...,
casting: _CastingSafe = ...,
optimize: _OptimizeKind = ...,
) -> Any: ...
@overload
def einsum(
__subscripts: str,
*operands: _ArrayLikeUInt_co,
out: None = ...,
dtype: Optional[_DTypeLikeUInt] = ...,
order: _OrderKACF = ...,
casting: _CastingSafe = ...,
optimize: _OptimizeKind = ...,
) -> Any: ...
@overload
def einsum(
__subscripts: str,
*operands: _ArrayLikeInt_co,
out: None = ...,
dtype: Optional[_DTypeLikeInt] = ...,
order: _OrderKACF = ...,
casting: _CastingSafe = ...,
optimize: _OptimizeKind = ...,
) -> Any: ...
@overload
def einsum(
__subscripts: str,
*operands: _ArrayLikeFloat_co,
out: None = ...,
dtype: Optional[_DTypeLikeFloat] = ...,
order: _OrderKACF = ...,
casting: _CastingSafe = ...,
optimize: _OptimizeKind = ...,
) -> Any: ...
@overload
def einsum(
__subscripts: str,
*operands: _ArrayLikeComplex_co,
out: None = ...,
dtype: Optional[_DTypeLikeComplex] = ...,
order: _OrderKACF = ...,
casting: _CastingSafe = ...,
optimize: _OptimizeKind = ...,
) -> Any: ...
@overload
def einsum(
__subscripts: str,
*operands: Any,
casting: _CastingUnsafe,
dtype: Optional[_DTypeLikeComplex_co] = ...,
out: None = ...,
order: _OrderKACF = ...,
optimize: _OptimizeKind = ...,
) -> Any: ...
@overload
def einsum(
__subscripts: str,
*operands: _ArrayLikeComplex_co,
out: _ArrayType,
dtype: Optional[_DTypeLikeComplex_co] = ...,
order: _OrderKACF = ...,
casting: _CastingSafe = ...,
optimize: _OptimizeKind = ...,
) -> _ArrayType: ...
@overload
def einsum(
__subscripts: str,
*operands: Any,
out: _ArrayType,
casting: _CastingUnsafe,
dtype: Optional[_DTypeLikeComplex_co] = ...,
order: _OrderKACF = ...,
optimize: _OptimizeKind = ...,
) -> _ArrayType: ...
# NOTE: `einsum_call` is a hidden kwarg unavailable for public use.
# It is therefore excluded from the signatures below.
# NOTE: In practice the list consists of a `str` (first element)
# and a variable number of integer tuples.
def einsum_path(
__subscripts: str,
*operands: _ArrayLikeComplex_co,
optimize: _OptimizeKind = ...,
) -> Tuple[List[Any], str]: ...
|