linalg.pyi 7.3 KB

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