aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/numpy/core/einsumfunc.pyi
diff options
context:
space:
mode:
Diffstat (limited to '.venv/lib/python3.12/site-packages/numpy/core/einsumfunc.pyi')
-rw-r--r--.venv/lib/python3.12/site-packages/numpy/core/einsumfunc.pyi187
1 files changed, 187 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/numpy/core/einsumfunc.pyi b/.venv/lib/python3.12/site-packages/numpy/core/einsumfunc.pyi
new file mode 100644
index 00000000..ad483bb9
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/numpy/core/einsumfunc.pyi
@@ -0,0 +1,187 @@
+from collections.abc import Sequence
+from typing import TypeVar, Any, overload, Union, Literal
+
+from numpy import (
+ ndarray,
+ dtype,
+ bool_,
+ number,
+ _OrderKACF,
+)
+from numpy._typing import (
+ _ArrayLikeBool_co,
+ _ArrayLikeUInt_co,
+ _ArrayLikeInt_co,
+ _ArrayLikeFloat_co,
+ _ArrayLikeComplex_co,
+ _ArrayLikeObject_co,
+ _DTypeLikeBool,
+ _DTypeLikeUInt,
+ _DTypeLikeInt,
+ _DTypeLikeFloat,
+ _DTypeLikeComplex,
+ _DTypeLikeComplex_co,
+ _DTypeLikeObject,
+)
+
+_ArrayType = TypeVar(
+ "_ArrayType",
+ bound=ndarray[Any, dtype[Union[bool_, number[Any]]]],
+)
+
+_OptimizeKind = 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 easily be done with
+# a mypy plugin.
+# Something like `is_scalar = bool(__subscripts.partition("->")[-1])`
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeBool_co,
+ out: None = ...,
+ dtype: None | _DTypeLikeBool = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeUInt_co,
+ out: None = ...,
+ dtype: None | _DTypeLikeUInt = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeInt_co,
+ out: None = ...,
+ dtype: None | _DTypeLikeInt = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeFloat_co,
+ out: None = ...,
+ dtype: None | _DTypeLikeFloat = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeComplex_co,
+ out: None = ...,
+ dtype: None | _DTypeLikeComplex = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: Any,
+ casting: _CastingUnsafe,
+ dtype: None | _DTypeLikeComplex_co = ...,
+ out: None = ...,
+ order: _OrderKACF = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeComplex_co,
+ out: _ArrayType,
+ dtype: None | _DTypeLikeComplex_co = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> _ArrayType: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: Any,
+ out: _ArrayType,
+ casting: _CastingUnsafe,
+ dtype: None | _DTypeLikeComplex_co = ...,
+ order: _OrderKACF = ...,
+ optimize: _OptimizeKind = ...,
+) -> _ArrayType: ...
+
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeObject_co,
+ out: None = ...,
+ dtype: None | _DTypeLikeObject = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: Any,
+ casting: _CastingUnsafe,
+ dtype: None | _DTypeLikeObject = ...,
+ out: None = ...,
+ order: _OrderKACF = ...,
+ optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeObject_co,
+ out: _ArrayType,
+ dtype: None | _DTypeLikeObject = ...,
+ order: _OrderKACF = ...,
+ casting: _CastingSafe = ...,
+ optimize: _OptimizeKind = ...,
+) -> _ArrayType: ...
+@overload
+def einsum(
+ subscripts: str | _ArrayLikeInt_co,
+ /,
+ *operands: Any,
+ out: _ArrayType,
+ casting: _CastingUnsafe,
+ dtype: None | _DTypeLikeObject = ...,
+ 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 | _ArrayLikeInt_co,
+ /,
+ *operands: _ArrayLikeComplex_co | _DTypeLikeObject,
+ optimize: _OptimizeKind = ...,
+) -> tuple[list[Any], str]: ...