numeric.pyi 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660
  1. from collections.abc import Callable, Sequence
  2. from typing import (
  3. Any,
  4. overload,
  5. TypeVar,
  6. Literal,
  7. SupportsAbs,
  8. SupportsIndex,
  9. NoReturn,
  10. )
  11. if sys.version_info >= (3, 10):
  12. from typing import TypeGuard
  13. else:
  14. from typing_extensions import TypeGuard
  15. from numpy import (
  16. ComplexWarning as ComplexWarning,
  17. generic,
  18. unsignedinteger,
  19. signedinteger,
  20. floating,
  21. complexfloating,
  22. bool_,
  23. int_,
  24. intp,
  25. float64,
  26. timedelta64,
  27. object_,
  28. _OrderKACF,
  29. _OrderCF,
  30. )
  31. from numpy._typing import (
  32. ArrayLike,
  33. NDArray,
  34. DTypeLike,
  35. _ShapeLike,
  36. _DTypeLike,
  37. _ArrayLike,
  38. _SupportsArrayFunc,
  39. _ScalarLike_co,
  40. _ArrayLikeBool_co,
  41. _ArrayLikeUInt_co,
  42. _ArrayLikeInt_co,
  43. _ArrayLikeFloat_co,
  44. _ArrayLikeComplex_co,
  45. _ArrayLikeTD64_co,
  46. _ArrayLikeObject_co,
  47. _ArrayLikeUnknown,
  48. )
  49. _T = TypeVar("_T")
  50. _SCT = TypeVar("_SCT", bound=generic)
  51. _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
  52. _CorrelateMode = Literal["valid", "same", "full"]
  53. __all__: list[str]
  54. @overload
  55. def zeros_like(
  56. a: _ArrayType,
  57. dtype: None = ...,
  58. order: _OrderKACF = ...,
  59. subok: Literal[True] = ...,
  60. shape: None = ...,
  61. ) -> _ArrayType: ...
  62. @overload
  63. def zeros_like(
  64. a: _ArrayLike[_SCT],
  65. dtype: None = ...,
  66. order: _OrderKACF = ...,
  67. subok: bool = ...,
  68. shape: None | _ShapeLike = ...,
  69. ) -> NDArray[_SCT]: ...
  70. @overload
  71. def zeros_like(
  72. a: object,
  73. dtype: None = ...,
  74. order: _OrderKACF = ...,
  75. subok: bool = ...,
  76. shape: None | _ShapeLike= ...,
  77. ) -> NDArray[Any]: ...
  78. @overload
  79. def zeros_like(
  80. a: Any,
  81. dtype: _DTypeLike[_SCT],
  82. order: _OrderKACF = ...,
  83. subok: bool = ...,
  84. shape: None | _ShapeLike= ...,
  85. ) -> NDArray[_SCT]: ...
  86. @overload
  87. def zeros_like(
  88. a: Any,
  89. dtype: DTypeLike,
  90. order: _OrderKACF = ...,
  91. subok: bool = ...,
  92. shape: None | _ShapeLike= ...,
  93. ) -> NDArray[Any]: ...
  94. @overload
  95. def ones(
  96. shape: _ShapeLike,
  97. dtype: None = ...,
  98. order: _OrderCF = ...,
  99. *,
  100. like: _SupportsArrayFunc = ...,
  101. ) -> NDArray[float64]: ...
  102. @overload
  103. def ones(
  104. shape: _ShapeLike,
  105. dtype: _DTypeLike[_SCT],
  106. order: _OrderCF = ...,
  107. *,
  108. like: _SupportsArrayFunc = ...,
  109. ) -> NDArray[_SCT]: ...
  110. @overload
  111. def ones(
  112. shape: _ShapeLike,
  113. dtype: DTypeLike,
  114. order: _OrderCF = ...,
  115. *,
  116. like: _SupportsArrayFunc = ...,
  117. ) -> NDArray[Any]: ...
  118. @overload
  119. def ones_like(
  120. a: _ArrayType,
  121. dtype: None = ...,
  122. order: _OrderKACF = ...,
  123. subok: Literal[True] = ...,
  124. shape: None = ...,
  125. ) -> _ArrayType: ...
  126. @overload
  127. def ones_like(
  128. a: _ArrayLike[_SCT],
  129. dtype: None = ...,
  130. order: _OrderKACF = ...,
  131. subok: bool = ...,
  132. shape: None | _ShapeLike = ...,
  133. ) -> NDArray[_SCT]: ...
  134. @overload
  135. def ones_like(
  136. a: object,
  137. dtype: None = ...,
  138. order: _OrderKACF = ...,
  139. subok: bool = ...,
  140. shape: None | _ShapeLike= ...,
  141. ) -> NDArray[Any]: ...
  142. @overload
  143. def ones_like(
  144. a: Any,
  145. dtype: _DTypeLike[_SCT],
  146. order: _OrderKACF = ...,
  147. subok: bool = ...,
  148. shape: None | _ShapeLike= ...,
  149. ) -> NDArray[_SCT]: ...
  150. @overload
  151. def ones_like(
  152. a: Any,
  153. dtype: DTypeLike,
  154. order: _OrderKACF = ...,
  155. subok: bool = ...,
  156. shape: None | _ShapeLike= ...,
  157. ) -> NDArray[Any]: ...
  158. @overload
  159. def full(
  160. shape: _ShapeLike,
  161. fill_value: Any,
  162. dtype: None = ...,
  163. order: _OrderCF = ...,
  164. *,
  165. like: _SupportsArrayFunc = ...,
  166. ) -> NDArray[Any]: ...
  167. @overload
  168. def full(
  169. shape: _ShapeLike,
  170. fill_value: Any,
  171. dtype: _DTypeLike[_SCT],
  172. order: _OrderCF = ...,
  173. *,
  174. like: _SupportsArrayFunc = ...,
  175. ) -> NDArray[_SCT]: ...
  176. @overload
  177. def full(
  178. shape: _ShapeLike,
  179. fill_value: Any,
  180. dtype: DTypeLike,
  181. order: _OrderCF = ...,
  182. *,
  183. like: _SupportsArrayFunc = ...,
  184. ) -> NDArray[Any]: ...
  185. @overload
  186. def full_like(
  187. a: _ArrayType,
  188. fill_value: Any,
  189. dtype: None = ...,
  190. order: _OrderKACF = ...,
  191. subok: Literal[True] = ...,
  192. shape: None = ...,
  193. ) -> _ArrayType: ...
  194. @overload
  195. def full_like(
  196. a: _ArrayLike[_SCT],
  197. fill_value: Any,
  198. dtype: None = ...,
  199. order: _OrderKACF = ...,
  200. subok: bool = ...,
  201. shape: None | _ShapeLike = ...,
  202. ) -> NDArray[_SCT]: ...
  203. @overload
  204. def full_like(
  205. a: object,
  206. fill_value: Any,
  207. dtype: None = ...,
  208. order: _OrderKACF = ...,
  209. subok: bool = ...,
  210. shape: None | _ShapeLike= ...,
  211. ) -> NDArray[Any]: ...
  212. @overload
  213. def full_like(
  214. a: Any,
  215. fill_value: Any,
  216. dtype: _DTypeLike[_SCT],
  217. order: _OrderKACF = ...,
  218. subok: bool = ...,
  219. shape: None | _ShapeLike= ...,
  220. ) -> NDArray[_SCT]: ...
  221. @overload
  222. def full_like(
  223. a: Any,
  224. fill_value: Any,
  225. dtype: DTypeLike,
  226. order: _OrderKACF = ...,
  227. subok: bool = ...,
  228. shape: None | _ShapeLike= ...,
  229. ) -> NDArray[Any]: ...
  230. @overload
  231. def count_nonzero(
  232. a: ArrayLike,
  233. axis: None = ...,
  234. *,
  235. keepdims: Literal[False] = ...,
  236. ) -> int: ...
  237. @overload
  238. def count_nonzero(
  239. a: ArrayLike,
  240. axis: _ShapeLike = ...,
  241. *,
  242. keepdims: bool = ...,
  243. ) -> Any: ... # TODO: np.intp or ndarray[np.intp]
  244. def isfortran(a: NDArray[Any] | generic) -> bool: ...
  245. def argwhere(a: ArrayLike) -> NDArray[intp]: ...
  246. def flatnonzero(a: ArrayLike) -> NDArray[intp]: ...
  247. @overload
  248. def correlate(
  249. a: _ArrayLikeUnknown,
  250. v: _ArrayLikeUnknown,
  251. mode: _CorrelateMode = ...,
  252. ) -> NDArray[Any]: ...
  253. @overload
  254. def correlate(
  255. a: _ArrayLikeBool_co,
  256. v: _ArrayLikeBool_co,
  257. mode: _CorrelateMode = ...,
  258. ) -> NDArray[bool_]: ...
  259. @overload
  260. def correlate(
  261. a: _ArrayLikeUInt_co,
  262. v: _ArrayLikeUInt_co,
  263. mode: _CorrelateMode = ...,
  264. ) -> NDArray[unsignedinteger[Any]]: ...
  265. @overload
  266. def correlate(
  267. a: _ArrayLikeInt_co,
  268. v: _ArrayLikeInt_co,
  269. mode: _CorrelateMode = ...,
  270. ) -> NDArray[signedinteger[Any]]: ...
  271. @overload
  272. def correlate(
  273. a: _ArrayLikeFloat_co,
  274. v: _ArrayLikeFloat_co,
  275. mode: _CorrelateMode = ...,
  276. ) -> NDArray[floating[Any]]: ...
  277. @overload
  278. def correlate(
  279. a: _ArrayLikeComplex_co,
  280. v: _ArrayLikeComplex_co,
  281. mode: _CorrelateMode = ...,
  282. ) -> NDArray[complexfloating[Any, Any]]: ...
  283. @overload
  284. def correlate(
  285. a: _ArrayLikeTD64_co,
  286. v: _ArrayLikeTD64_co,
  287. mode: _CorrelateMode = ...,
  288. ) -> NDArray[timedelta64]: ...
  289. @overload
  290. def correlate(
  291. a: _ArrayLikeObject_co,
  292. v: _ArrayLikeObject_co,
  293. mode: _CorrelateMode = ...,
  294. ) -> NDArray[object_]: ...
  295. @overload
  296. def convolve(
  297. a: _ArrayLikeUnknown,
  298. v: _ArrayLikeUnknown,
  299. mode: _CorrelateMode = ...,
  300. ) -> NDArray[Any]: ...
  301. @overload
  302. def convolve(
  303. a: _ArrayLikeBool_co,
  304. v: _ArrayLikeBool_co,
  305. mode: _CorrelateMode = ...,
  306. ) -> NDArray[bool_]: ...
  307. @overload
  308. def convolve(
  309. a: _ArrayLikeUInt_co,
  310. v: _ArrayLikeUInt_co,
  311. mode: _CorrelateMode = ...,
  312. ) -> NDArray[unsignedinteger[Any]]: ...
  313. @overload
  314. def convolve(
  315. a: _ArrayLikeInt_co,
  316. v: _ArrayLikeInt_co,
  317. mode: _CorrelateMode = ...,
  318. ) -> NDArray[signedinteger[Any]]: ...
  319. @overload
  320. def convolve(
  321. a: _ArrayLikeFloat_co,
  322. v: _ArrayLikeFloat_co,
  323. mode: _CorrelateMode = ...,
  324. ) -> NDArray[floating[Any]]: ...
  325. @overload
  326. def convolve(
  327. a: _ArrayLikeComplex_co,
  328. v: _ArrayLikeComplex_co,
  329. mode: _CorrelateMode = ...,
  330. ) -> NDArray[complexfloating[Any, Any]]: ...
  331. @overload
  332. def convolve(
  333. a: _ArrayLikeTD64_co,
  334. v: _ArrayLikeTD64_co,
  335. mode: _CorrelateMode = ...,
  336. ) -> NDArray[timedelta64]: ...
  337. @overload
  338. def convolve(
  339. a: _ArrayLikeObject_co,
  340. v: _ArrayLikeObject_co,
  341. mode: _CorrelateMode = ...,
  342. ) -> NDArray[object_]: ...
  343. @overload
  344. def outer(
  345. a: _ArrayLikeUnknown,
  346. b: _ArrayLikeUnknown,
  347. out: None = ...,
  348. ) -> NDArray[Any]: ...
  349. @overload
  350. def outer(
  351. a: _ArrayLikeBool_co,
  352. b: _ArrayLikeBool_co,
  353. out: None = ...,
  354. ) -> NDArray[bool_]: ...
  355. @overload
  356. def outer(
  357. a: _ArrayLikeUInt_co,
  358. b: _ArrayLikeUInt_co,
  359. out: None = ...,
  360. ) -> NDArray[unsignedinteger[Any]]: ...
  361. @overload
  362. def outer(
  363. a: _ArrayLikeInt_co,
  364. b: _ArrayLikeInt_co,
  365. out: None = ...,
  366. ) -> NDArray[signedinteger[Any]]: ...
  367. @overload
  368. def outer(
  369. a: _ArrayLikeFloat_co,
  370. b: _ArrayLikeFloat_co,
  371. out: None = ...,
  372. ) -> NDArray[floating[Any]]: ...
  373. @overload
  374. def outer(
  375. a: _ArrayLikeComplex_co,
  376. b: _ArrayLikeComplex_co,
  377. out: None = ...,
  378. ) -> NDArray[complexfloating[Any, Any]]: ...
  379. @overload
  380. def outer(
  381. a: _ArrayLikeTD64_co,
  382. b: _ArrayLikeTD64_co,
  383. out: None = ...,
  384. ) -> NDArray[timedelta64]: ...
  385. @overload
  386. def outer(
  387. a: _ArrayLikeObject_co,
  388. b: _ArrayLikeObject_co,
  389. out: None = ...,
  390. ) -> NDArray[object_]: ...
  391. @overload
  392. def outer(
  393. a: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
  394. b: _ArrayLikeComplex_co | _ArrayLikeTD64_co | _ArrayLikeObject_co,
  395. out: _ArrayType,
  396. ) -> _ArrayType: ...
  397. @overload
  398. def tensordot(
  399. a: _ArrayLikeUnknown,
  400. b: _ArrayLikeUnknown,
  401. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  402. ) -> NDArray[Any]: ...
  403. @overload
  404. def tensordot(
  405. a: _ArrayLikeBool_co,
  406. b: _ArrayLikeBool_co,
  407. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  408. ) -> NDArray[bool_]: ...
  409. @overload
  410. def tensordot(
  411. a: _ArrayLikeUInt_co,
  412. b: _ArrayLikeUInt_co,
  413. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  414. ) -> NDArray[unsignedinteger[Any]]: ...
  415. @overload
  416. def tensordot(
  417. a: _ArrayLikeInt_co,
  418. b: _ArrayLikeInt_co,
  419. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  420. ) -> NDArray[signedinteger[Any]]: ...
  421. @overload
  422. def tensordot(
  423. a: _ArrayLikeFloat_co,
  424. b: _ArrayLikeFloat_co,
  425. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  426. ) -> NDArray[floating[Any]]: ...
  427. @overload
  428. def tensordot(
  429. a: _ArrayLikeComplex_co,
  430. b: _ArrayLikeComplex_co,
  431. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  432. ) -> NDArray[complexfloating[Any, Any]]: ...
  433. @overload
  434. def tensordot(
  435. a: _ArrayLikeTD64_co,
  436. b: _ArrayLikeTD64_co,
  437. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  438. ) -> NDArray[timedelta64]: ...
  439. @overload
  440. def tensordot(
  441. a: _ArrayLikeObject_co,
  442. b: _ArrayLikeObject_co,
  443. axes: int | tuple[_ShapeLike, _ShapeLike] = ...,
  444. ) -> NDArray[object_]: ...
  445. @overload
  446. def roll(
  447. a: _ArrayLike[_SCT],
  448. shift: _ShapeLike,
  449. axis: None | _ShapeLike = ...,
  450. ) -> NDArray[_SCT]: ...
  451. @overload
  452. def roll(
  453. a: ArrayLike,
  454. shift: _ShapeLike,
  455. axis: None | _ShapeLike = ...,
  456. ) -> NDArray[Any]: ...
  457. def rollaxis(
  458. a: NDArray[_SCT],
  459. axis: int,
  460. start: int = ...,
  461. ) -> NDArray[_SCT]: ...
  462. def moveaxis(
  463. a: NDArray[_SCT],
  464. source: _ShapeLike,
  465. destination: _ShapeLike,
  466. ) -> NDArray[_SCT]: ...
  467. @overload
  468. def cross(
  469. a: _ArrayLikeUnknown,
  470. b: _ArrayLikeUnknown,
  471. axisa: int = ...,
  472. axisb: int = ...,
  473. axisc: int = ...,
  474. axis: None | int = ...,
  475. ) -> NDArray[Any]: ...
  476. @overload
  477. def cross(
  478. a: _ArrayLikeBool_co,
  479. b: _ArrayLikeBool_co,
  480. axisa: int = ...,
  481. axisb: int = ...,
  482. axisc: int = ...,
  483. axis: None | int = ...,
  484. ) -> NoReturn: ...
  485. @overload
  486. def cross(
  487. a: _ArrayLikeUInt_co,
  488. b: _ArrayLikeUInt_co,
  489. axisa: int = ...,
  490. axisb: int = ...,
  491. axisc: int = ...,
  492. axis: None | int = ...,
  493. ) -> NDArray[unsignedinteger[Any]]: ...
  494. @overload
  495. def cross(
  496. a: _ArrayLikeInt_co,
  497. b: _ArrayLikeInt_co,
  498. axisa: int = ...,
  499. axisb: int = ...,
  500. axisc: int = ...,
  501. axis: None | int = ...,
  502. ) -> NDArray[signedinteger[Any]]: ...
  503. @overload
  504. def cross(
  505. a: _ArrayLikeFloat_co,
  506. b: _ArrayLikeFloat_co,
  507. axisa: int = ...,
  508. axisb: int = ...,
  509. axisc: int = ...,
  510. axis: None | int = ...,
  511. ) -> NDArray[floating[Any]]: ...
  512. @overload
  513. def cross(
  514. a: _ArrayLikeComplex_co,
  515. b: _ArrayLikeComplex_co,
  516. axisa: int = ...,
  517. axisb: int = ...,
  518. axisc: int = ...,
  519. axis: None | int = ...,
  520. ) -> NDArray[complexfloating[Any, Any]]: ...
  521. @overload
  522. def cross(
  523. a: _ArrayLikeObject_co,
  524. b: _ArrayLikeObject_co,
  525. axisa: int = ...,
  526. axisb: int = ...,
  527. axisc: int = ...,
  528. axis: None | int = ...,
  529. ) -> NDArray[object_]: ...
  530. @overload
  531. def indices(
  532. dimensions: Sequence[int],
  533. dtype: type[int] = ...,
  534. sparse: Literal[False] = ...,
  535. ) -> NDArray[int_]: ...
  536. @overload
  537. def indices(
  538. dimensions: Sequence[int],
  539. dtype: type[int] = ...,
  540. sparse: Literal[True] = ...,
  541. ) -> tuple[NDArray[int_], ...]: ...
  542. @overload
  543. def indices(
  544. dimensions: Sequence[int],
  545. dtype: _DTypeLike[_SCT],
  546. sparse: Literal[False] = ...,
  547. ) -> NDArray[_SCT]: ...
  548. @overload
  549. def indices(
  550. dimensions: Sequence[int],
  551. dtype: _DTypeLike[_SCT],
  552. sparse: Literal[True],
  553. ) -> tuple[NDArray[_SCT], ...]: ...
  554. @overload
  555. def indices(
  556. dimensions: Sequence[int],
  557. dtype: DTypeLike,
  558. sparse: Literal[False] = ...,
  559. ) -> NDArray[Any]: ...
  560. @overload
  561. def indices(
  562. dimensions: Sequence[int],
  563. dtype: DTypeLike,
  564. sparse: Literal[True],
  565. ) -> tuple[NDArray[Any], ...]: ...
  566. def fromfunction(
  567. function: Callable[..., _T],
  568. shape: Sequence[int],
  569. *,
  570. dtype: DTypeLike = ...,
  571. like: _SupportsArrayFunc = ...,
  572. **kwargs: Any,
  573. ) -> _T: ...
  574. def isscalar(element: object) -> TypeGuard[
  575. generic | bool | int | float | complex | str | bytes | memoryview
  576. ]: ...
  577. def binary_repr(num: SupportsIndex, width: None | int = ...) -> str: ...
  578. def base_repr(
  579. number: SupportsAbs[float],
  580. base: float = ...,
  581. padding: SupportsIndex = ...,
  582. ) -> str: ...
  583. @overload
  584. def identity(
  585. n: int,
  586. dtype: None = ...,
  587. *,
  588. like: _SupportsArrayFunc = ...,
  589. ) -> NDArray[float64]: ...
  590. @overload
  591. def identity(
  592. n: int,
  593. dtype: _DTypeLike[_SCT],
  594. *,
  595. like: _SupportsArrayFunc = ...,
  596. ) -> NDArray[_SCT]: ...
  597. @overload
  598. def identity(
  599. n: int,
  600. dtype: DTypeLike,
  601. *,
  602. like: _SupportsArrayFunc = ...,
  603. ) -> NDArray[Any]: ...
  604. def allclose(
  605. a: ArrayLike,
  606. b: ArrayLike,
  607. rtol: float = ...,
  608. atol: float = ...,
  609. equal_nan: bool = ...,
  610. ) -> bool: ...
  611. @overload
  612. def isclose(
  613. a: _ScalarLike_co,
  614. b: _ScalarLike_co,
  615. rtol: float = ...,
  616. atol: float = ...,
  617. equal_nan: bool = ...,
  618. ) -> bool_: ...
  619. @overload
  620. def isclose(
  621. a: ArrayLike,
  622. b: ArrayLike,
  623. rtol: float = ...,
  624. atol: float = ...,
  625. equal_nan: bool = ...,
  626. ) -> NDArray[bool_]: ...
  627. def array_equal(a1: ArrayLike, a2: ArrayLike, equal_nan: bool = ...) -> bool: ...
  628. def array_equiv(a1: ArrayLike, a2: ArrayLike) -> bool: ...