diff options
| author | Bas van Beek <bas.vanbeek@hotmail.com> | 2023-05-08 15:31:33 +0200 |
|---|---|---|
| committer | Bas van Beek <bas.vanbeek@hotmail.com> | 2023-05-08 15:31:33 +0200 |
| commit | cc99207f78c84a7c7a89c5c1523371d1dcfe23bf (patch) | |
| tree | 0c09c483f49c709cb9b22541df07805d0e4d7393 /numpy/core | |
| parent | 89f00a66af2fd73391caa5ce5aae41ba3d693b1a (diff) | |
| download | numpy-cc99207f78c84a7c7a89c5c1523371d1dcfe23bf.tar.gz | |
TYP: Let `np.einsum` accept `object` dtypes
xref https://github.com/numpy/numpy/pull/18053
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]: ... |