about summary refs log tree commit diff
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]: ...