diff options
author | Bas van Beek <b.f.van.beek@vu.nl> | 2021-09-18 12:49:57 +0200 |
---|---|---|
committer | Bas van Beek <b.f.van.beek@vu.nl> | 2021-09-18 12:49:57 +0200 |
commit | f8958a2305c88f1c97b2cd5256d7e29faca93f13 (patch) | |
tree | 4c73aa61847f9f665daffeee0f7b50c5ee498cf9 /numpy | |
parent | d5cee25a90fa651c9e865c72de872b9a4a0a8016 (diff) | |
download | numpy-f8958a2305c88f1c97b2cd5256d7e29faca93f13.tar.gz |
MAINT: Let `eig` and `eigvals` return a union for real inputs
Whether an array of a `floating` or `complexfloating` dtype is returned depends on the values within in the supplied matrix.
Co-Authored-By: Kevin Sheppard <bashtage@users.noreply.github.com>
Diffstat (limited to 'numpy')
-rw-r--r-- | numpy/linalg/linalg.pyi | 18 | ||||
-rw-r--r-- | numpy/typing/tests/data/reveal/linalg.py | 8 |
2 files changed, 17 insertions, 9 deletions
diff --git a/numpy/linalg/linalg.pyi b/numpy/linalg/linalg.pyi index cd9fd1493..3ce0e2229 100644 --- a/numpy/linalg/linalg.pyi +++ b/numpy/linalg/linalg.pyi @@ -10,7 +10,15 @@ from typing import ( Tuple, ) -from numpy import generic, int32, floating, float64, complexfloating +from numpy import ( + generic, + floating, + complexfloating, + int32, + float64, + complex128, +) + from numpy.typing import ( NDArray, ArrayLike, @@ -109,9 +117,9 @@ def qr(a: _ArrayLikeFloat_co, mode: _ModeKind = ...) -> _2Tuple[NDArray[floating def qr(a: _ArrayLikeComplex_co, mode: _ModeKind = ...) -> _2Tuple[NDArray[complexfloating[Any, Any]]]: ... @overload -def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64]: ... +def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64] | NDArray[complex128]: ... @overload -def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... +def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]] | NDArray[complexfloating[Any, Any]]: ... @overload def eigvals(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... @@ -121,9 +129,9 @@ def eigvalsh(a: _ArrayLikeInt_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[ def eigvalsh(a: _ArrayLikeComplex_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[floating[Any]]: ... @overload -def eig(a: _ArrayLikeInt_co) -> _2Tuple[NDArray[float64]]: ... +def eig(a: _ArrayLikeInt_co) -> _2Tuple[NDArray[float64]] | _2Tuple[NDArray[complex128]]: ... @overload -def eig(a: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]]: ... +def eig(a: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]] | _2Tuple[NDArray[complexfloating[Any, Any]]]: ... @overload def eig(a: _ArrayLikeComplex_co) -> _2Tuple[NDArray[complexfloating[Any, Any]]]: ... diff --git a/numpy/typing/tests/data/reveal/linalg.py b/numpy/typing/tests/data/reveal/linalg.py index f4f52d085..fecdc0d37 100644 --- a/numpy/typing/tests/data/reveal/linalg.py +++ b/numpy/typing/tests/data/reveal/linalg.py @@ -37,16 +37,16 @@ reveal_type(np.linalg.qr(AR_i8)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{flo reveal_type(np.linalg.qr(AR_f8)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]] reveal_type(np.linalg.qr(AR_c16)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]] -reveal_type(np.linalg.eigvals(AR_i8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] -reveal_type(np.linalg.eigvals(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] +reveal_type(np.linalg.eigvals(AR_i8)) # E: Union[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{complex128}]]] +reveal_type(np.linalg.eigvals(AR_f8)) # E: Union[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]] reveal_type(np.linalg.eigvals(AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]] reveal_type(np.linalg.eigvalsh(AR_i8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] reveal_type(np.linalg.eigvalsh(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] reveal_type(np.linalg.eigvalsh(AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] -reveal_type(np.linalg.eig(AR_i8)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{float64}]]] -reveal_type(np.linalg.eig(AR_f8)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]] +reveal_type(np.linalg.eig(AR_i8)) # E: Union[Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{float64}]]], Tuple[numpy.ndarray[Any, numpy.dtype[{complex128}]], numpy.ndarray[Any, numpy.dtype[{complex128}]]]] +reveal_type(np.linalg.eig(AR_f8)) # E: Union[Tuple[numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]], Tuple[numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]]] reveal_type(np.linalg.eig(AR_c16)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]] reveal_type(np.linalg.eigh(AR_i8)) # E: Tuple[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{float64}]]] |