from collections.abc import Iterable
from typing import (
Literal as L,
overload,
TypeAlias,
TypeVar,
Any,
SupportsIndex,
SupportsInt,
NamedTuple,
Generic,
)
import numpy as np
from numpy import (
generic,
floating,
complexfloating,
signedinteger,
unsignedinteger,
timedelta64,
object_,
int32,
float64,
complex128,
)
from numpy.linalg import LinAlgError as LinAlgError
from numpy._typing import (
NDArray,
ArrayLike,
_ArrayLikeUnknown,
_ArrayLikeBool_co,
_ArrayLikeInt_co,
_ArrayLikeUInt_co,
_ArrayLikeFloat_co,
_ArrayLikeComplex_co,
_ArrayLikeTD64_co,
_ArrayLikeObject_co,
DTypeLike,
)
_T = TypeVar("_T")
_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
_SCT = TypeVar("_SCT", bound=generic, covariant=True)
_SCT2 = TypeVar("_SCT2", bound=generic, covariant=True)
_2Tuple: TypeAlias = tuple[_T, _T]
_ModeKind: TypeAlias = L["reduced", "complete", "r", "raw"]
__all__: list[str]
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]]: ...