|
from collections.abc import Iterable |
|
from typing import ( |
|
Literal as L, |
|
overload, |
|
TypeAlias, |
|
TypeVar, |
|
Any, |
|
SupportsIndex, |
|
SupportsInt, |
|
NamedTuple, |
|
) |
|
|
|
import numpy as np |
|
from numpy import ( |
|
# re-exports |
|
vecdot, |
|
|
|
# other |
|
generic, |
|
floating, |
|
complexfloating, |
|
signedinteger, |
|
unsignedinteger, |
|
timedelta64, |
|
object_, |
|
int32, |
|
float64, |
|
complex128, |
|
) |
|
from numpy.linalg import LinAlgError |
|
from numpy._core.fromnumeric import matrix_transpose |
|
from numpy._core.numeric import tensordot |
|
from numpy._typing import ( |
|
NDArray, |
|
ArrayLike, |
|
DTypeLike, |
|
_ArrayLikeUnknown, |
|
_ArrayLikeBool_co, |
|
_ArrayLikeInt_co, |
|
_ArrayLikeUInt_co, |
|
_ArrayLikeFloat_co, |
|
_ArrayLikeComplex_co, |
|
_ArrayLikeTD64_co, |
|
_ArrayLikeObject_co, |
|
) |
|
|
|
__all__ = [ |
|
"matrix_power", |
|
"solve", |
|
"tensorsolve", |
|
"tensorinv", |
|
"inv", |
|
"cholesky", |
|
"eigvals", |
|
"eigvalsh", |
|
"pinv", |
|
"slogdet", |
|
"det", |
|
"svd", |
|
"svdvals", |
|
"eig", |
|
"eigh", |
|
"lstsq", |
|
"norm", |
|
"qr", |
|
"cond", |
|
"matrix_rank", |
|
"LinAlgError", |
|
"multi_dot", |
|
"trace", |
|
"diagonal", |
|
"cross", |
|
"outer", |
|
"tensordot", |
|
"matmul", |
|
"matrix_transpose", |
|
"matrix_norm", |
|
"vector_norm", |
|
"vecdot", |
|
] |
|
|
|
_T = TypeVar("_T") |
|
_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any]) |
|
_SCT2 = TypeVar("_SCT2", bound=generic, covariant=True) |
|
|
|
_2Tuple: TypeAlias = tuple[_T, _T] |
|
_ModeKind: TypeAlias = L["reduced", "complete", "r", "raw"] |
|
|
|
|
|
class EigResult(NamedTuple): |
|
eigenvalues: NDArray[Any] |
|
eigenvectors: NDArray[Any] |
|
|
|
class EighResult(NamedTuple): |
|
eigenvalues: NDArray[Any] |
|
eigenvectors: NDArray[Any] |
|
|
|
class QRResult(NamedTuple): |
|
Q: NDArray[Any] |
|
R: NDArray[Any] |
|
|
|
class SlogdetResult(NamedTuple): |
|
# TODO: `sign` and `logabsdet` are scalars for input 2D arrays and |
|
# a `(x.ndim - 2)`` dimensionl arrays otherwise |
|
sign: Any |
|
logabsdet: Any |
|
|
|
class SVDResult(NamedTuple): |
|
U: NDArray[Any] |
|
S: NDArray[Any] |
|
Vh: NDArray[Any] |
|
|
|
@overload |
|
def tensorsolve( |
|
a: _ArrayLikeInt_co, |
|
b: _ArrayLikeInt_co, |
|
axes: None | Iterable[int] =..., |
|
) -> NDArray[float64]: ... |
|
@overload |
|
def tensorsolve( |
|
a: _ArrayLikeFloat_co, |
|
b: _ArrayLikeFloat_co, |
|
axes: None | Iterable[int] =..., |
|
) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def tensorsolve( |
|
a: _ArrayLikeComplex_co, |
|
b: _ArrayLikeComplex_co, |
|
axes: None | Iterable[int] =..., |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
@overload |
|
def solve( |
|
a: _ArrayLikeInt_co, |
|
b: _ArrayLikeInt_co, |
|
) -> NDArray[float64]: ... |
|
@overload |
|
def solve( |
|
a: _ArrayLikeFloat_co, |
|
b: _ArrayLikeFloat_co, |
|
) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def solve( |
|
a: _ArrayLikeComplex_co, |
|
b: _ArrayLikeComplex_co, |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
@overload |
|
def tensorinv( |
|
a: _ArrayLikeInt_co, |
|
ind: int = ..., |
|
) -> NDArray[float64]: ... |
|
@overload |
|
def tensorinv( |
|
a: _ArrayLikeFloat_co, |
|
ind: int = ..., |
|
) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def tensorinv( |
|
a: _ArrayLikeComplex_co, |
|
ind: int = ..., |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
@overload |
|
def inv(a: _ArrayLikeInt_co) -> NDArray[float64]: ... |
|
@overload |
|
def inv(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def inv(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
# TODO: The supported input and output dtypes are dependent on the value of `n`. |
|
# For example: `n < 0` always casts integer types to float64 |
|
def matrix_power( |
|
a: _ArrayLikeComplex_co | _ArrayLikeObject_co, |
|
n: SupportsIndex, |
|
) -> NDArray[Any]: ... |
|
|
|
@overload |
|
def cholesky(a: _ArrayLikeInt_co) -> NDArray[float64]: ... |
|
@overload |
|
def cholesky(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def cholesky(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
@overload |
|
def outer(x1: _ArrayLikeUnknown, x2: _ArrayLikeUnknown) -> NDArray[Any]: ... |
|
@overload |
|
def outer(x1: _ArrayLikeBool_co, x2: _ArrayLikeBool_co) -> NDArray[np.bool]: ... |
|
@overload |
|
def outer(x1: _ArrayLikeUInt_co, x2: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... |
|
@overload |
|
def outer(x1: _ArrayLikeInt_co, x2: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... |
|
@overload |
|
def outer(x1: _ArrayLikeFloat_co, x2: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def outer( |
|
x1: _ArrayLikeComplex_co, |
|
x2: _ArrayLikeComplex_co, |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
@overload |
|
def outer( |
|
x1: _ArrayLikeTD64_co, |
|
x2: _ArrayLikeTD64_co, |
|
out: None = ..., |
|
) -> NDArray[timedelta64]: ... |
|
@overload |
|
def outer(x1: _ArrayLikeObject_co, x2: _ArrayLikeObject_co) -> NDArray[object_]: ... |
|
@overload |
|
def outer( |
|
x1: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, |
|
x2: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co, |
|
) -> _ArrayType: ... |
|
|
|
@overload |
|
def qr(a: _ArrayLikeInt_co, mode: _ModeKind = ...) -> QRResult: ... |
|
@overload |
|
def qr(a: _ArrayLikeFloat_co, mode: _ModeKind = ...) -> QRResult: ... |
|
@overload |
|
def qr(a: _ArrayLikeComplex_co, mode: _ModeKind = ...) -> QRResult: ... |
|
|
|
@overload |
|
def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64] | NDArray[complex128]: ... |
|
@overload |
|
def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]] | NDArray[complexfloating[Any, Any]]: ... |
|
@overload |
|
def eigvals(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
@overload |
|
def eigvalsh(a: _ArrayLikeInt_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[float64]: ... |
|
@overload |
|
def eigvalsh(a: _ArrayLikeComplex_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[floating[Any]]: ... |
|
|
|
@overload |
|
def eig(a: _ArrayLikeInt_co) -> EigResult: ... |
|
@overload |
|
def eig(a: _ArrayLikeFloat_co) -> EigResult: ... |
|
@overload |
|
def eig(a: _ArrayLikeComplex_co) -> EigResult: ... |
|
|
|
@overload |
|
def eigh( |
|
a: _ArrayLikeInt_co, |
|
UPLO: L["L", "U", "l", "u"] = ..., |
|
) -> EighResult: ... |
|
@overload |
|
def eigh( |
|
a: _ArrayLikeFloat_co, |
|
UPLO: L["L", "U", "l", "u"] = ..., |
|
) -> EighResult: ... |
|
@overload |
|
def eigh( |
|
a: _ArrayLikeComplex_co, |
|
UPLO: L["L", "U", "l", "u"] = ..., |
|
) -> EighResult: ... |
|
|
|
@overload |
|
def svd( |
|
a: _ArrayLikeInt_co, |
|
full_matrices: bool = ..., |
|
compute_uv: L[True] = ..., |
|
hermitian: bool = ..., |
|
) -> SVDResult: ... |
|
@overload |
|
def svd( |
|
a: _ArrayLikeFloat_co, |
|
full_matrices: bool = ..., |
|
compute_uv: L[True] = ..., |
|
hermitian: bool = ..., |
|
) -> SVDResult: ... |
|
@overload |
|
def svd( |
|
a: _ArrayLikeComplex_co, |
|
full_matrices: bool = ..., |
|
compute_uv: L[True] = ..., |
|
hermitian: bool = ..., |
|
) -> SVDResult: ... |
|
@overload |
|
def svd( |
|
a: _ArrayLikeInt_co, |
|
full_matrices: bool = ..., |
|
compute_uv: L[False] = ..., |
|
hermitian: bool = ..., |
|
) -> NDArray[float64]: ... |
|
@overload |
|
def svd( |
|
a: _ArrayLikeComplex_co, |
|
full_matrices: bool = ..., |
|
compute_uv: L[False] = ..., |
|
hermitian: bool = ..., |
|
) -> NDArray[floating[Any]]: ... |
|
|
|
def svdvals( |
|
x: _ArrayLikeInt_co | _ArrayLikeFloat_co | _ArrayLikeComplex_co |
|
) -> NDArray[floating[Any]]: ... |
|
|
|
# TODO: Returns a scalar for 2D arrays and |
|
# a `(x.ndim - 2)`` dimensionl array otherwise |
|
def cond(x: _ArrayLikeComplex_co, p: None | float | L["fro", "nuc"] = ...) -> Any: ... |
|
|
|
# TODO: Returns `int` for <2D arrays and `intp` otherwise |
|
def matrix_rank( |
|
A: _ArrayLikeComplex_co, |
|
tol: None | _ArrayLikeFloat_co = ..., |
|
hermitian: bool = ..., |
|
*, |
|
rtol: None | _ArrayLikeFloat_co = ..., |
|
) -> Any: ... |
|
|
|
@overload |
|
def pinv( |
|
a: _ArrayLikeInt_co, |
|
rcond: _ArrayLikeFloat_co = ..., |
|
hermitian: bool = ..., |
|
) -> NDArray[float64]: ... |
|
@overload |
|
def pinv( |
|
a: _ArrayLikeFloat_co, |
|
rcond: _ArrayLikeFloat_co = ..., |
|
hermitian: bool = ..., |
|
) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def pinv( |
|
a: _ArrayLikeComplex_co, |
|
rcond: _ArrayLikeFloat_co = ..., |
|
hermitian: bool = ..., |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
# TODO: Returns a 2-tuple of scalars for 2D arrays and |
|
# a 2-tuple of `(a.ndim - 2)`` dimensionl arrays otherwise |
|
def slogdet(a: _ArrayLikeComplex_co) -> SlogdetResult: ... |
|
|
|
# TODO: Returns a 2-tuple of scalars for 2D arrays and |
|
# a 2-tuple of `(a.ndim - 2)`` dimensionl arrays otherwise |
|
def det(a: _ArrayLikeComplex_co) -> Any: ... |
|
|
|
@overload |
|
def lstsq(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co, rcond: None | float = ...) -> tuple[ |
|
NDArray[float64], |
|
NDArray[float64], |
|
int32, |
|
NDArray[float64], |
|
]: ... |
|
@overload |
|
def lstsq(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, rcond: None | float = ...) -> tuple[ |
|
NDArray[floating[Any]], |
|
NDArray[floating[Any]], |
|
int32, |
|
NDArray[floating[Any]], |
|
]: ... |
|
@overload |
|
def lstsq(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, rcond: None | float = ...) -> tuple[ |
|
NDArray[complexfloating[Any, Any]], |
|
NDArray[floating[Any]], |
|
int32, |
|
NDArray[floating[Any]], |
|
]: ... |
|
|
|
@overload |
|
def norm( |
|
x: ArrayLike, |
|
ord: None | float | L["fro", "nuc"] = ..., |
|
axis: None = ..., |
|
keepdims: bool = ..., |
|
) -> floating[Any]: ... |
|
@overload |
|
def norm( |
|
x: ArrayLike, |
|
ord: None | float | L["fro", "nuc"] = ..., |
|
axis: SupportsInt | SupportsIndex | tuple[int, ...] = ..., |
|
keepdims: bool = ..., |
|
) -> Any: ... |
|
|
|
@overload |
|
def matrix_norm( |
|
x: ArrayLike, |
|
ord: None | float | L["fro", "nuc"] = ..., |
|
keepdims: bool = ..., |
|
) -> floating[Any]: ... |
|
@overload |
|
def matrix_norm( |
|
x: ArrayLike, |
|
ord: None | float | L["fro", "nuc"] = ..., |
|
keepdims: bool = ..., |
|
) -> Any: ... |
|
|
|
@overload |
|
def vector_norm( |
|
x: ArrayLike, |
|
axis: None = ..., |
|
ord: None | float = ..., |
|
keepdims: bool = ..., |
|
) -> floating[Any]: ... |
|
@overload |
|
def vector_norm( |
|
x: ArrayLike, |
|
axis: SupportsInt | SupportsIndex | tuple[int, ...] = ..., |
|
ord: None | float = ..., |
|
keepdims: bool = ..., |
|
) -> Any: ... |
|
|
|
# TODO: Returns a scalar or array |
|
def multi_dot( |
|
arrays: Iterable[_ArrayLikeComplex_co | _ArrayLikeObject_co | _ArrayLikeTD64_co], |
|
*, |
|
out: None | NDArray[Any] = ..., |
|
) -> Any: ... |
|
|
|
def diagonal( |
|
x: ArrayLike, # >= 2D array |
|
offset: SupportsIndex = ..., |
|
) -> NDArray[Any]: ... |
|
|
|
def trace( |
|
x: ArrayLike, # >= 2D array |
|
offset: SupportsIndex = ..., |
|
dtype: DTypeLike = ..., |
|
) -> Any: ... |
|
|
|
@overload |
|
def cross( |
|
a: _ArrayLikeUInt_co, |
|
b: _ArrayLikeUInt_co, |
|
axis: int = ..., |
|
) -> NDArray[unsignedinteger[Any]]: ... |
|
@overload |
|
def cross( |
|
a: _ArrayLikeInt_co, |
|
b: _ArrayLikeInt_co, |
|
axis: int = ..., |
|
) -> NDArray[signedinteger[Any]]: ... |
|
@overload |
|
def cross( |
|
a: _ArrayLikeFloat_co, |
|
b: _ArrayLikeFloat_co, |
|
axis: int = ..., |
|
) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def cross( |
|
a: _ArrayLikeComplex_co, |
|
b: _ArrayLikeComplex_co, |
|
axis: int = ..., |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
|
|
@overload |
|
def matmul( |
|
x1: _ArrayLikeInt_co, |
|
x2: _ArrayLikeInt_co, |
|
) -> NDArray[signedinteger[Any]]: ... |
|
@overload |
|
def matmul( |
|
x1: _ArrayLikeUInt_co, |
|
x2: _ArrayLikeUInt_co, |
|
) -> NDArray[unsignedinteger[Any]]: ... |
|
@overload |
|
def matmul( |
|
x1: _ArrayLikeFloat_co, |
|
x2: _ArrayLikeFloat_co, |
|
) -> NDArray[floating[Any]]: ... |
|
@overload |
|
def matmul( |
|
x1: _ArrayLikeComplex_co, |
|
x2: _ArrayLikeComplex_co, |
|
) -> NDArray[complexfloating[Any, Any]]: ... |
|
|