diff options
Diffstat (limited to 'numpy/core')
| -rw-r--r-- | numpy/core/einsumfunc.pyi | 53 |
1 files changed, 48 insertions, 5 deletions
diff --git a/numpy/core/einsumfunc.pyi b/numpy/core/einsumfunc.pyi index c811a5783..ad483bb90 100644 --- a/numpy/core/einsumfunc.pyi +++ b/numpy/core/einsumfunc.pyi @@ -5,10 +5,6 @@ from numpy import ( ndarray, dtype, bool_, - unsignedinteger, - signedinteger, - floating, - complexfloating, number, _OrderKACF, ) @@ -18,12 +14,14 @@ from numpy._typing import ( _ArrayLikeInt_co, _ArrayLikeFloat_co, _ArrayLikeComplex_co, + _ArrayLikeObject_co, _DTypeLikeBool, _DTypeLikeUInt, _DTypeLikeInt, _DTypeLikeFloat, _DTypeLikeComplex, _DTypeLikeComplex_co, + _DTypeLikeObject, ) _ArrayType = TypeVar( @@ -132,6 +130,51 @@ def einsum( 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) @@ -139,6 +182,6 @@ def einsum( def einsum_path( subscripts: str | _ArrayLikeInt_co, /, - *operands: _ArrayLikeComplex_co, + *operands: _ArrayLikeComplex_co | _DTypeLikeObject, optimize: _OptimizeKind = ..., ) -> tuple[list[Any], str]: ... |