numeric.pyi 14 KB

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