from collections.abc import Iterable from typing import ( Any, Literal as L, NamedTuple, Never, SupportsIndex, SupportsInt, TypeAlias, TypeVar, overload, ) import numpy as np from numpy import ( complex128, complexfloating, float64, floating, int32, object_, signedinteger, timedelta64, unsignedinteger, vecdot, ) from numpy._core.fromnumeric import matrix_transpose from numpy._globals import _NoValueType from numpy._typing import ( ArrayLike, DTypeLike, NDArray, _ArrayLike, _ArrayLikeBool_co, _ArrayLikeComplex_co, _ArrayLikeFloat_co, _ArrayLikeInt_co, _ArrayLikeNumber_co, _ArrayLikeObject_co, _ArrayLikeTD64_co, _ArrayLikeUInt_co, _NestedSequence, _ShapeLike, ) from numpy.linalg import LinAlgError __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", ] _NumberT = TypeVar("_NumberT", bound=np.number) _NumericScalarT = TypeVar("_NumericScalarT", bound=np.number | np.timedelta64 | np.object_) _ModeKind: TypeAlias = L["reduced", "complete", "r", "raw"] ### fortran_int = np.intc 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: Iterable[int] | None = None, ) -> NDArray[float64]: ... @overload def tensorsolve( a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, axes: Iterable[int] | None = None, ) -> NDArray[floating]: ... @overload def tensorsolve( a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, axes: Iterable[int] | None = None, ) -> NDArray[complexfloating]: ... @overload def solve( a: _ArrayLikeInt_co, b: _ArrayLikeInt_co, ) -> NDArray[float64]: ... @overload def solve( a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, ) -> NDArray[floating]: ... @overload def solve( a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, ) -> NDArray[complexfloating]: ... @overload def tensorinv( a: _ArrayLikeInt_co, ind: int = 2, ) -> NDArray[float64]: ... @overload def tensorinv( a: _ArrayLikeFloat_co, ind: int = 2, ) -> NDArray[floating]: ... @overload def tensorinv( a: _ArrayLikeComplex_co, ind: int = 2, ) -> NDArray[complexfloating]: ... @overload def inv(a: _ArrayLikeInt_co) -> NDArray[float64]: ... @overload def inv(a: _ArrayLikeFloat_co) -> NDArray[floating]: ... @overload def inv(a: _ArrayLikeComplex_co) -> NDArray[complexfloating]: ... # 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, /, *, upper: bool = False) -> NDArray[float64]: ... @overload def cholesky(a: _ArrayLikeFloat_co, /, *, upper: bool = False) -> NDArray[floating]: ... @overload def cholesky(a: _ArrayLikeComplex_co, /, *, upper: bool = False) -> NDArray[complexfloating]: ... @overload def outer(x1: _ArrayLike[Never], x2: _ArrayLike[Never], /) -> NDArray[Any]: ... @overload def outer(x1: _ArrayLikeBool_co, x2: _ArrayLikeBool_co, /) -> NDArray[np.bool]: ... @overload def outer(x1: _ArrayLike[_NumberT], x2: _ArrayLike[_NumberT], /) -> NDArray[_NumberT]: ... @overload def outer(x1: _ArrayLikeUInt_co, x2: _ArrayLikeUInt_co, /) -> NDArray[unsignedinteger]: ... @overload def outer(x1: _ArrayLikeInt_co, x2: _ArrayLikeInt_co, /) -> NDArray[signedinteger]: ... @overload def outer(x1: _ArrayLikeFloat_co, x2: _ArrayLikeFloat_co, /) -> NDArray[floating]: ... @overload def outer(x1: _ArrayLikeComplex_co, x2: _ArrayLikeComplex_co, /) -> NDArray[complexfloating]: ... @overload def outer(x1: _ArrayLikeTD64_co, x2: _ArrayLikeTD64_co, /) -> 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, /, ) -> NDArray[Any]: ... @overload def qr(a: _ArrayLikeInt_co, mode: _ModeKind = "reduced") -> QRResult: ... @overload def qr(a: _ArrayLikeFloat_co, mode: _ModeKind = "reduced") -> QRResult: ... @overload def qr(a: _ArrayLikeComplex_co, mode: _ModeKind = "reduced") -> QRResult: ... @overload def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64] | NDArray[complex128]: ... @overload def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating] | NDArray[complexfloating]: ... @overload def eigvals(a: _ArrayLikeComplex_co) -> NDArray[complexfloating]: ... @overload def eigvalsh(a: _ArrayLikeInt_co, UPLO: L["L", "U", "l", "u"] = "L") -> NDArray[float64]: ... @overload def eigvalsh(a: _ArrayLikeComplex_co, UPLO: L["L", "U", "l", "u"] = "L") -> NDArray[floating]: ... @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"] = "L", ) -> EighResult: ... @overload def eigh( a: _ArrayLikeFloat_co, UPLO: L["L", "U", "l", "u"] = "L", ) -> EighResult: ... @overload def eigh( a: _ArrayLikeComplex_co, UPLO: L["L", "U", "l", "u"] = "L", ) -> EighResult: ... @overload def svd( a: _ArrayLikeInt_co, full_matrices: bool = True, compute_uv: L[True] = True, hermitian: bool = False, ) -> SVDResult: ... @overload def svd( a: _ArrayLikeFloat_co, full_matrices: bool = True, compute_uv: L[True] = True, hermitian: bool = False, ) -> SVDResult: ... @overload def svd( a: _ArrayLikeComplex_co, full_matrices: bool = True, compute_uv: L[True] = True, hermitian: bool = False, ) -> SVDResult: ... @overload def svd( a: _ArrayLikeInt_co, full_matrices: bool = True, *, compute_uv: L[False], hermitian: bool = False, ) -> NDArray[float64]: ... @overload def svd( a: _ArrayLikeInt_co, full_matrices: bool, compute_uv: L[False], hermitian: bool = False, ) -> NDArray[float64]: ... @overload def svd( a: _ArrayLikeComplex_co, full_matrices: bool = True, *, compute_uv: L[False], hermitian: bool = False, ) -> NDArray[floating]: ... @overload def svd( a: _ArrayLikeComplex_co, full_matrices: bool, compute_uv: L[False], hermitian: bool = False, ) -> NDArray[floating]: ... # the ignored `overload-overlap` mypy error below is a false-positive @overload def svdvals( # type: ignore[overload-overlap] x: _ArrayLike[np.float64 | np.complex128 | np.integer | np.bool] | _NestedSequence[complex], / ) -> NDArray[np.float64]: ... @overload def svdvals(x: _ArrayLike[np.float32 | np.complex64], /) -> NDArray[np.float32]: ... @overload def svdvals(x: _ArrayLikeNumber_co, /) -> NDArray[floating]: ... # TODO: Returns a scalar for 2D arrays and # a `(x.ndim - 2)`` dimensionl array otherwise def cond(x: _ArrayLikeComplex_co, p: float | L["fro", "nuc"] | None = None) -> Any: ... # TODO: Returns `int` for <2D arrays and `intp` otherwise def matrix_rank( A: _ArrayLikeComplex_co, tol: _ArrayLikeFloat_co | None = None, hermitian: bool = False, *, rtol: _ArrayLikeFloat_co | None = None, ) -> Any: ... @overload def pinv( a: _ArrayLikeInt_co, rcond: _ArrayLikeFloat_co | None = None, hermitian: bool = False, *, rtol: _ArrayLikeFloat_co | _NoValueType = ..., ) -> NDArray[float64]: ... @overload def pinv( a: _ArrayLikeFloat_co, rcond: _ArrayLikeFloat_co | None = None, hermitian: bool = False, *, rtol: _ArrayLikeFloat_co | _NoValueType = ..., ) -> NDArray[floating]: ... @overload def pinv( a: _ArrayLikeComplex_co, rcond: _ArrayLikeFloat_co | None = None, hermitian: bool = False, *, rtol: _ArrayLikeFloat_co | _NoValueType = ..., ) -> NDArray[complexfloating]: ... # 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: float | None = None) -> tuple[ NDArray[float64], NDArray[float64], int32, NDArray[float64], ]: ... @overload def lstsq(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, rcond: float | None = None) -> tuple[ NDArray[floating], NDArray[floating], int32, NDArray[floating], ]: ... @overload def lstsq(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, rcond: float | None = None) -> tuple[ NDArray[complexfloating], NDArray[floating], int32, NDArray[floating], ]: ... @overload def norm( x: ArrayLike, ord: float | L["fro", "nuc"] | None = None, axis: None = None, keepdims: L[False] = False, ) -> floating: ... @overload def norm( x: ArrayLike, ord: float | L["fro", "nuc"] | None, axis: SupportsInt | SupportsIndex | tuple[int, ...] | None, keepdims: bool = False, ) -> Any: ... @overload def norm( x: ArrayLike, ord: float | L["fro", "nuc"] | None = None, *, axis: SupportsInt | SupportsIndex | tuple[int, ...] | None, keepdims: bool = False, ) -> Any: ... @overload def matrix_norm( x: ArrayLike, /, *, ord: float | L["fro", "nuc"] | None = "fro", keepdims: L[False] = False, ) -> floating: ... @overload def matrix_norm( x: ArrayLike, /, *, ord: float | L["fro", "nuc"] | None = "fro", keepdims: bool = False, ) -> Any: ... @overload def vector_norm( x: ArrayLike, /, *, axis: None = None, ord: float | None = 2, keepdims: L[False] = False, ) -> floating: ... @overload def vector_norm( x: ArrayLike, /, *, axis: SupportsInt | SupportsIndex | tuple[int, ...], ord: float | None = 2, keepdims: bool = False, ) -> Any: ... # keep in sync with numpy._core.numeric.tensordot (ignoring `/, *`) @overload def tensordot( a: _ArrayLike[_NumericScalarT], b: _ArrayLike[_NumericScalarT], /, *, axes: int | tuple[_ShapeLike, _ShapeLike] = 2, ) -> NDArray[_NumericScalarT]: ... @overload def tensordot( a: _ArrayLikeBool_co, b: _ArrayLikeBool_co, /, *, axes: int | tuple[_ShapeLike, _ShapeLike] = 2, ) -> NDArray[np.bool_]: ... @overload def tensordot( a: _ArrayLikeInt_co, b: _ArrayLikeInt_co, /, *, axes: int | tuple[_ShapeLike, _ShapeLike] = 2, ) -> NDArray[np.int_ | Any]: ... @overload def tensordot( a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, /, *, axes: int | tuple[_ShapeLike, _ShapeLike] = 2, ) -> NDArray[np.float64 | Any]: ... @overload def tensordot( a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, /, *, axes: int | tuple[_ShapeLike, _ShapeLike] = 2, ) -> NDArray[np.complex128 | Any]: ... # TODO: Returns a scalar or array def multi_dot( arrays: Iterable[_ArrayLikeComplex_co | _ArrayLikeObject_co | _ArrayLikeTD64_co], *, out: NDArray[Any] | None = None, ) -> Any: ... def diagonal( x: ArrayLike, # >= 2D array /, *, offset: SupportsIndex = 0, ) -> NDArray[Any]: ... def trace( x: ArrayLike, # >= 2D array /, *, offset: SupportsIndex = 0, dtype: DTypeLike | None = None, ) -> Any: ... @overload def cross( x1: _ArrayLikeUInt_co, x2: _ArrayLikeUInt_co, /, *, axis: int = -1, ) -> NDArray[unsignedinteger]: ... @overload def cross( x1: _ArrayLikeInt_co, x2: _ArrayLikeInt_co, /, *, axis: int = -1, ) -> NDArray[signedinteger]: ... @overload def cross( x1: _ArrayLikeFloat_co, x2: _ArrayLikeFloat_co, /, *, axis: int = -1, ) -> NDArray[floating]: ... @overload def cross( x1: _ArrayLikeComplex_co, x2: _ArrayLikeComplex_co, /, *, axis: int = -1, ) -> NDArray[complexfloating]: ... @overload def matmul(x1: _ArrayLike[_NumberT], x2: _ArrayLike[_NumberT], /) -> NDArray[_NumberT]: ... @overload def matmul(x1: _ArrayLikeUInt_co, x2: _ArrayLikeUInt_co, /) -> NDArray[unsignedinteger]: ... @overload def matmul(x1: _ArrayLikeInt_co, x2: _ArrayLikeInt_co, /) -> NDArray[signedinteger]: ... @overload def matmul(x1: _ArrayLikeFloat_co, x2: _ArrayLikeFloat_co, /) -> NDArray[floating]: ... @overload def matmul(x1: _ArrayLikeComplex_co, x2: _ArrayLikeComplex_co, /) -> NDArray[complexfloating]: ...