Skip to content

Commit ff80d91

Browse files
author
Bas van Beek
committed
ENH: Add annotations for np.linalg 
1 parent 1d584d6 commit ff80d91

File tree

1 file changed

+274
-20
lines changed

1 file changed

+274
-20
lines changed

‎numpy/linalg/linalg.pyi‎

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

0 commit comments

Comments
 (0)