linalg.pyi 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297
  1. from collections.abc import Iterable
  2. from typing import (
  3. Literal as L,
  4. overload,
  5. TypeVar,
  6. Any,
  7. SupportsIndex,
  8. SupportsInt,
  9. NamedTuple,
  10. Generic,
  11. )
  12. from numpy import (
  13. generic,
  14. floating,
  15. complexfloating,
  16. int32,
  17. float64,
  18. complex128,
  19. )
  20. from numpy.linalg import LinAlgError as LinAlgError
  21. from numpy._typing import (
  22. NDArray,
  23. ArrayLike,
  24. _ArrayLikeInt_co,
  25. _ArrayLikeFloat_co,
  26. _ArrayLikeComplex_co,
  27. _ArrayLikeTD64_co,
  28. _ArrayLikeObject_co,
  29. )
  30. _T = TypeVar("_T")
  31. _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
  32. _SCT = TypeVar("_SCT", bound=generic, covariant=True)
  33. _SCT2 = TypeVar("_SCT2", bound=generic, covariant=True)
  34. _2Tuple = tuple[_T, _T]
  35. _ModeKind = L["reduced", "complete", "r", "raw"]
  36. __all__: list[str]
  37. class EigResult(NamedTuple):
  38. eigenvalues: NDArray[Any]
  39. eigenvectors: NDArray[Any]
  40. class EighResult(NamedTuple):
  41. eigenvalues: NDArray[Any]
  42. eigenvectors: NDArray[Any]
  43. class QRResult(NamedTuple):
  44. Q: NDArray[Any]
  45. R: NDArray[Any]
  46. class SlogdetResult(NamedTuple):
  47. # TODO: `sign` and `logabsdet` are scalars for input 2D arrays and
  48. # a `(x.ndim - 2)`` dimensionl arrays otherwise
  49. sign: Any
  50. logabsdet: Any
  51. class SVDResult(NamedTuple):
  52. U: NDArray[Any]
  53. S: NDArray[Any]
  54. Vh: NDArray[Any]
  55. @overload
  56. def tensorsolve(
  57. a: _ArrayLikeInt_co,
  58. b: _ArrayLikeInt_co,
  59. axes: None | Iterable[int] =...,
  60. ) -> NDArray[float64]: ...
  61. @overload
  62. def tensorsolve(
  63. a: _ArrayLikeFloat_co,
  64. b: _ArrayLikeFloat_co,
  65. axes: None | Iterable[int] =...,
  66. ) -> NDArray[floating[Any]]: ...
  67. @overload
  68. def tensorsolve(
  69. a: _ArrayLikeComplex_co,
  70. b: _ArrayLikeComplex_co,
  71. axes: None | Iterable[int] =...,
  72. ) -> NDArray[complexfloating[Any, Any]]: ...
  73. @overload
  74. def solve(
  75. a: _ArrayLikeInt_co,
  76. b: _ArrayLikeInt_co,
  77. ) -> NDArray[float64]: ...
  78. @overload
  79. def solve(
  80. a: _ArrayLikeFloat_co,
  81. b: _ArrayLikeFloat_co,
  82. ) -> NDArray[floating[Any]]: ...
  83. @overload
  84. def solve(
  85. a: _ArrayLikeComplex_co,
  86. b: _ArrayLikeComplex_co,
  87. ) -> NDArray[complexfloating[Any, Any]]: ...
  88. @overload
  89. def tensorinv(
  90. a: _ArrayLikeInt_co,
  91. ind: int = ...,
  92. ) -> NDArray[float64]: ...
  93. @overload
  94. def tensorinv(
  95. a: _ArrayLikeFloat_co,
  96. ind: int = ...,
  97. ) -> NDArray[floating[Any]]: ...
  98. @overload
  99. def tensorinv(
  100. a: _ArrayLikeComplex_co,
  101. ind: int = ...,
  102. ) -> NDArray[complexfloating[Any, Any]]: ...
  103. @overload
  104. def inv(a: _ArrayLikeInt_co) -> NDArray[float64]: ...
  105. @overload
  106. def inv(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ...
  107. @overload
  108. def inv(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...
  109. # TODO: The supported input and output dtypes are dependent on the value of `n`.
  110. # For example: `n < 0` always casts integer types to float64
  111. def matrix_power(
  112. a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
  113. n: SupportsIndex,
  114. ) -> NDArray[Any]: ...
  115. @overload
  116. def cholesky(a: _ArrayLikeInt_co) -> NDArray[float64]: ...
  117. @overload
  118. def cholesky(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ...
  119. @overload
  120. def cholesky(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...
  121. @overload
  122. def qr(a: _ArrayLikeInt_co, mode: _ModeKind = ...) -> QRResult: ...
  123. @overload
  124. def qr(a: _ArrayLikeFloat_co, mode: _ModeKind = ...) -> QRResult: ...
  125. @overload
  126. def qr(a: _ArrayLikeComplex_co, mode: _ModeKind = ...) -> QRResult: ...
  127. @overload
  128. def eigvals(a: _ArrayLikeInt_co) -> NDArray[float64] | NDArray[complex128]: ...
  129. @overload
  130. def eigvals(a: _ArrayLikeFloat_co) -> NDArray[floating[Any]] | NDArray[complexfloating[Any, Any]]: ...
  131. @overload
  132. def eigvals(a: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...
  133. @overload
  134. def eigvalsh(a: _ArrayLikeInt_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[float64]: ...
  135. @overload
  136. def eigvalsh(a: _ArrayLikeComplex_co, UPLO: L["L", "U", "l", "u"] = ...) -> NDArray[floating[Any]]: ...
  137. @overload
  138. def eig(a: _ArrayLikeInt_co) -> EigResult: ...
  139. @overload
  140. def eig(a: _ArrayLikeFloat_co) -> EigResult: ...
  141. @overload
  142. def eig(a: _ArrayLikeComplex_co) -> EigResult: ...
  143. @overload
  144. def eigh(
  145. a: _ArrayLikeInt_co,
  146. UPLO: L["L", "U", "l", "u"] = ...,
  147. ) -> EighResult: ...
  148. @overload
  149. def eigh(
  150. a: _ArrayLikeFloat_co,
  151. UPLO: L["L", "U", "l", "u"] = ...,
  152. ) -> EighResult: ...
  153. @overload
  154. def eigh(
  155. a: _ArrayLikeComplex_co,
  156. UPLO: L["L", "U", "l", "u"] = ...,
  157. ) -> EighResult: ...
  158. @overload
  159. def svd(
  160. a: _ArrayLikeInt_co,
  161. full_matrices: bool = ...,
  162. compute_uv: L[True] = ...,
  163. hermitian: bool = ...,
  164. ) -> SVDResult: ...
  165. @overload
  166. def svd(
  167. a: _ArrayLikeFloat_co,
  168. full_matrices: bool = ...,
  169. compute_uv: L[True] = ...,
  170. hermitian: bool = ...,
  171. ) -> SVDResult: ...
  172. @overload
  173. def svd(
  174. a: _ArrayLikeComplex_co,
  175. full_matrices: bool = ...,
  176. compute_uv: L[True] = ...,
  177. hermitian: bool = ...,
  178. ) -> SVDResult: ...
  179. @overload
  180. def svd(
  181. a: _ArrayLikeInt_co,
  182. full_matrices: bool = ...,
  183. compute_uv: L[False] = ...,
  184. hermitian: bool = ...,
  185. ) -> NDArray[float64]: ...
  186. @overload
  187. def svd(
  188. a: _ArrayLikeComplex_co,
  189. full_matrices: bool = ...,
  190. compute_uv: L[False] = ...,
  191. hermitian: bool = ...,
  192. ) -> NDArray[floating[Any]]: ...
  193. # TODO: Returns a scalar for 2D arrays and
  194. # a `(x.ndim - 2)`` dimensionl array otherwise
  195. def cond(x: _ArrayLikeComplex_co, p: None | float | L["fro", "nuc"] = ...) -> Any: ...
  196. # TODO: Returns `int` for <2D arrays and `intp` otherwise
  197. def matrix_rank(
  198. A: _ArrayLikeComplex_co,
  199. tol: None | _ArrayLikeFloat_co = ...,
  200. hermitian: bool = ...,
  201. ) -> Any: ...
  202. @overload
  203. def pinv(
  204. a: _ArrayLikeInt_co,
  205. rcond: _ArrayLikeFloat_co = ...,
  206. hermitian: bool = ...,
  207. ) -> NDArray[float64]: ...
  208. @overload
  209. def pinv(
  210. a: _ArrayLikeFloat_co,
  211. rcond: _ArrayLikeFloat_co = ...,
  212. hermitian: bool = ...,
  213. ) -> NDArray[floating[Any]]: ...
  214. @overload
  215. def pinv(
  216. a: _ArrayLikeComplex_co,
  217. rcond: _ArrayLikeFloat_co = ...,
  218. hermitian: bool = ...,
  219. ) -> NDArray[complexfloating[Any, Any]]: ...
  220. # TODO: Returns a 2-tuple of scalars for 2D arrays and
  221. # a 2-tuple of `(a.ndim - 2)`` dimensionl arrays otherwise
  222. def slogdet(a: _ArrayLikeComplex_co) -> SlogdetResult: ...
  223. # TODO: Returns a 2-tuple of scalars for 2D arrays and
  224. # a 2-tuple of `(a.ndim - 2)`` dimensionl arrays otherwise
  225. def det(a: _ArrayLikeComplex_co) -> Any: ...
  226. @overload
  227. def lstsq(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co, rcond: None | float = ...) -> tuple[
  228. NDArray[float64],
  229. NDArray[float64],
  230. int32,
  231. NDArray[float64],
  232. ]: ...
  233. @overload
  234. def lstsq(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co, rcond: None | float = ...) -> tuple[
  235. NDArray[floating[Any]],
  236. NDArray[floating[Any]],
  237. int32,
  238. NDArray[floating[Any]],
  239. ]: ...
  240. @overload
  241. def lstsq(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co, rcond: None | float = ...) -> tuple[
  242. NDArray[complexfloating[Any, Any]],
  243. NDArray[floating[Any]],
  244. int32,
  245. NDArray[floating[Any]],
  246. ]: ...
  247. @overload
  248. def norm(
  249. x: ArrayLike,
  250. ord: None | float | L["fro", "nuc"] = ...,
  251. axis: None = ...,
  252. keepdims: bool = ...,
  253. ) -> floating[Any]: ...
  254. @overload
  255. def norm(
  256. x: ArrayLike,
  257. ord: None | float | L["fro", "nuc"] = ...,
  258. axis: SupportsInt | SupportsIndex | tuple[int, ...] = ...,
  259. keepdims: bool = ...,
  260. ) -> Any: ...
  261. # TODO: Returns a scalar or array
  262. def multi_dot(
  263. arrays: Iterable[_ArrayLikeComplex_co | _ArrayLikeObject_co | _ArrayLikeTD64_co],
  264. *,
  265. out: None | NDArray[Any] = ...,
  266. ) -> Any: ...