_array_object.py 42 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118
  1. """
  2. Wrapper class around the ndarray object for the array API standard.
  3. The array API standard defines some behaviors differently than ndarray, in
  4. particular, type promotion rules are different (the standard has no
  5. value-based casting). The standard also specifies a more limited subset of
  6. array methods and functionalities than are implemented on ndarray. Since the
  7. goal of the array_api namespace is to be a minimal implementation of the array
  8. API standard, we need to define a separate wrapper class for the array_api
  9. namespace.
  10. The standard compliant class is only a wrapper class. It is *not* a subclass
  11. of ndarray.
  12. """
  13. from __future__ import annotations
  14. import operator
  15. from enum import IntEnum
  16. from ._creation_functions import asarray
  17. from ._dtypes import (
  18. _all_dtypes,
  19. _boolean_dtypes,
  20. _integer_dtypes,
  21. _integer_or_boolean_dtypes,
  22. _floating_dtypes,
  23. _numeric_dtypes,
  24. _result_type,
  25. _dtype_categories,
  26. )
  27. from typing import TYPE_CHECKING, Optional, Tuple, Union, Any, SupportsIndex
  28. import types
  29. if TYPE_CHECKING:
  30. from ._typing import Any, PyCapsule, Device, Dtype
  31. import numpy.typing as npt
  32. import numpy as np
  33. from numpy import array_api
  34. class Array:
  35. """
  36. n-d array object for the array API namespace.
  37. See the docstring of :py:obj:`np.ndarray <numpy.ndarray>` for more
  38. information.
  39. This is a wrapper around numpy.ndarray that restricts the usage to only
  40. those things that are required by the array API namespace. Note,
  41. attributes on this object that start with a single underscore are not part
  42. of the API specification and should only be used internally. This object
  43. should not be constructed directly. Rather, use one of the creation
  44. functions, such as asarray().
  45. """
  46. _array: np.ndarray
  47. # Use a custom constructor instead of __init__, as manually initializing
  48. # this class is not supported API.
  49. @classmethod
  50. def _new(cls, x, /):
  51. """
  52. This is a private method for initializing the array API Array
  53. object.
  54. Functions outside of the array_api submodule should not use this
  55. method. Use one of the creation functions instead, such as
  56. ``asarray``.
  57. """
  58. obj = super().__new__(cls)
  59. # Note: The spec does not have array scalars, only 0-D arrays.
  60. if isinstance(x, np.generic):
  61. # Convert the array scalar to a 0-D array
  62. x = np.asarray(x)
  63. if x.dtype not in _all_dtypes:
  64. raise TypeError(
  65. f"The array_api namespace does not support the dtype '{x.dtype}'"
  66. )
  67. obj._array = x
  68. return obj
  69. # Prevent Array() from working
  70. def __new__(cls, *args, **kwargs):
  71. raise TypeError(
  72. "The array_api Array object should not be instantiated directly. Use an array creation function, such as asarray(), instead."
  73. )
  74. # These functions are not required by the spec, but are implemented for
  75. # the sake of usability.
  76. def __str__(self: Array, /) -> str:
  77. """
  78. Performs the operation __str__.
  79. """
  80. return self._array.__str__().replace("array", "Array")
  81. def __repr__(self: Array, /) -> str:
  82. """
  83. Performs the operation __repr__.
  84. """
  85. suffix = f", dtype={self.dtype.name})"
  86. if 0 in self.shape:
  87. prefix = "empty("
  88. mid = str(self.shape)
  89. else:
  90. prefix = "Array("
  91. mid = np.array2string(self._array, separator=', ', prefix=prefix, suffix=suffix)
  92. return prefix + mid + suffix
  93. # This function is not required by the spec, but we implement it here for
  94. # convenience so that np.asarray(np.array_api.Array) will work.
  95. def __array__(self, dtype: None | np.dtype[Any] = None) -> npt.NDArray[Any]:
  96. """
  97. Warning: this method is NOT part of the array API spec. Implementers
  98. of other libraries need not include it, and users should not assume it
  99. will be present in other implementations.
  100. """
  101. return np.asarray(self._array, dtype=dtype)
  102. # These are various helper functions to make the array behavior match the
  103. # spec in places where it either deviates from or is more strict than
  104. # NumPy behavior
  105. def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
  106. """
  107. Helper function for operators to only allow specific input dtypes
  108. Use like
  109. other = self._check_allowed_dtypes(other, 'numeric', '__add__')
  110. if other is NotImplemented:
  111. return other
  112. """
  113. if self.dtype not in _dtype_categories[dtype_category]:
  114. raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
  115. if isinstance(other, (int, float, bool)):
  116. other = self._promote_scalar(other)
  117. elif isinstance(other, Array):
  118. if other.dtype not in _dtype_categories[dtype_category]:
  119. raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
  120. else:
  121. return NotImplemented
  122. # This will raise TypeError for type combinations that are not allowed
  123. # to promote in the spec (even if the NumPy array operator would
  124. # promote them).
  125. res_dtype = _result_type(self.dtype, other.dtype)
  126. if op.startswith("__i"):
  127. # Note: NumPy will allow in-place operators in some cases where
  128. # the type promoted operator does not match the left-hand side
  129. # operand. For example,
  130. # >>> a = np.array(1, dtype=np.int8)
  131. # >>> a += np.array(1, dtype=np.int16)
  132. # The spec explicitly disallows this.
  133. if res_dtype != self.dtype:
  134. raise TypeError(
  135. f"Cannot perform {op} with dtypes {self.dtype} and {other.dtype}"
  136. )
  137. return other
  138. # Helper function to match the type promotion rules in the spec
  139. def _promote_scalar(self, scalar):
  140. """
  141. Returns a promoted version of a Python scalar appropriate for use with
  142. operations on self.
  143. This may raise an OverflowError in cases where the scalar is an
  144. integer that is too large to fit in a NumPy integer dtype, or
  145. TypeError when the scalar type is incompatible with the dtype of self.
  146. """
  147. # Note: Only Python scalar types that match the array dtype are
  148. # allowed.
  149. if isinstance(scalar, bool):
  150. if self.dtype not in _boolean_dtypes:
  151. raise TypeError(
  152. "Python bool scalars can only be promoted with bool arrays"
  153. )
  154. elif isinstance(scalar, int):
  155. if self.dtype in _boolean_dtypes:
  156. raise TypeError(
  157. "Python int scalars cannot be promoted with bool arrays"
  158. )
  159. elif isinstance(scalar, float):
  160. if self.dtype not in _floating_dtypes:
  161. raise TypeError(
  162. "Python float scalars can only be promoted with floating-point arrays."
  163. )
  164. else:
  165. raise TypeError("'scalar' must be a Python scalar")
  166. # Note: scalars are unconditionally cast to the same dtype as the
  167. # array.
  168. # Note: the spec only specifies integer-dtype/int promotion
  169. # behavior for integers within the bounds of the integer dtype.
  170. # Outside of those bounds we use the default NumPy behavior (either
  171. # cast or raise OverflowError).
  172. return Array._new(np.array(scalar, self.dtype))
  173. @staticmethod
  174. def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
  175. """
  176. Normalize inputs to two arg functions to fix type promotion rules
  177. NumPy deviates from the spec type promotion rules in cases where one
  178. argument is 0-dimensional and the other is not. For example:
  179. >>> import numpy as np
  180. >>> a = np.array([1.0], dtype=np.float32)
  181. >>> b = np.array(1.0, dtype=np.float64)
  182. >>> np.add(a, b) # The spec says this should be float64
  183. array([2.], dtype=float32)
  184. To fix this, we add a dimension to the 0-dimension array before passing it
  185. through. This works because a dimension would be added anyway from
  186. broadcasting, so the resulting shape is the same, but this prevents NumPy
  187. from not promoting the dtype.
  188. """
  189. # Another option would be to use signature=(x1.dtype, x2.dtype, None),
  190. # but that only works for ufuncs, so we would have to call the ufuncs
  191. # directly in the operator methods. One should also note that this
  192. # sort of trick wouldn't work for functions like searchsorted, which
  193. # don't do normal broadcasting, but there aren't any functions like
  194. # that in the array API namespace.
  195. if x1.ndim == 0 and x2.ndim != 0:
  196. # The _array[None] workaround was chosen because it is relatively
  197. # performant. broadcast_to(x1._array, x2.shape) is much slower. We
  198. # could also manually type promote x2, but that is more complicated
  199. # and about the same performance as this.
  200. x1 = Array._new(x1._array[None])
  201. elif x2.ndim == 0 and x1.ndim != 0:
  202. x2 = Array._new(x2._array[None])
  203. return (x1, x2)
  204. # Note: A large fraction of allowed indices are disallowed here (see the
  205. # docstring below)
  206. def _validate_index(self, key):
  207. """
  208. Validate an index according to the array API.
  209. The array API specification only requires a subset of indices that are
  210. supported by NumPy. This function will reject any index that is
  211. allowed by NumPy but not required by the array API specification. We
  212. always raise ``IndexError`` on such indices (the spec does not require
  213. any specific behavior on them, but this makes the NumPy array API
  214. namespace a minimal implementation of the spec). See
  215. https://data-apis.org/array-api/latest/API_specification/indexing.html
  216. for the full list of required indexing behavior
  217. This function raises IndexError if the index ``key`` is invalid. It
  218. only raises ``IndexError`` on indices that are not already rejected by
  219. NumPy, as NumPy will already raise the appropriate error on such
  220. indices. ``shape`` may be None, in which case, only cases that are
  221. independent of the array shape are checked.
  222. The following cases are allowed by NumPy, but not specified by the array
  223. API specification:
  224. - Indices to not include an implicit ellipsis at the end. That is,
  225. every axis of an array must be explicitly indexed or an ellipsis
  226. included. This behaviour is sometimes referred to as flat indexing.
  227. - The start and stop of a slice may not be out of bounds. In
  228. particular, for a slice ``i:j:k`` on an axis of size ``n``, only the
  229. following are allowed:
  230. - ``i`` or ``j`` omitted (``None``).
  231. - ``-n <= i <= max(0, n - 1)``.
  232. - For ``k > 0`` or ``k`` omitted (``None``), ``-n <= j <= n``.
  233. - For ``k < 0``, ``-n - 1 <= j <= max(0, n - 1)``.
  234. - Boolean array indices are not allowed as part of a larger tuple
  235. index.
  236. - Integer array indices are not allowed (with the exception of 0-D
  237. arrays, which are treated the same as scalars).
  238. Additionally, it should be noted that indices that would return a
  239. scalar in NumPy will return a 0-D array. Array scalars are not allowed
  240. in the specification, only 0-D arrays. This is done in the
  241. ``Array._new`` constructor, not this function.
  242. """
  243. _key = key if isinstance(key, tuple) else (key,)
  244. for i in _key:
  245. if isinstance(i, bool) or not (
  246. isinstance(i, SupportsIndex) # i.e. ints
  247. or isinstance(i, slice)
  248. or i == Ellipsis
  249. or i is None
  250. or isinstance(i, Array)
  251. or isinstance(i, np.ndarray)
  252. ):
  253. raise IndexError(
  254. f"Single-axes index {i} has {type(i)=}, but only "
  255. "integers, slices (:), ellipsis (...), newaxis (None), "
  256. "zero-dimensional integer arrays and boolean arrays "
  257. "are specified in the Array API."
  258. )
  259. nonexpanding_key = []
  260. single_axes = []
  261. n_ellipsis = 0
  262. key_has_mask = False
  263. for i in _key:
  264. if i is not None:
  265. nonexpanding_key.append(i)
  266. if isinstance(i, Array) or isinstance(i, np.ndarray):
  267. if i.dtype in _boolean_dtypes:
  268. key_has_mask = True
  269. single_axes.append(i)
  270. else:
  271. # i must not be an array here, to avoid elementwise equals
  272. if i == Ellipsis:
  273. n_ellipsis += 1
  274. else:
  275. single_axes.append(i)
  276. n_single_axes = len(single_axes)
  277. if n_ellipsis > 1:
  278. return # handled by ndarray
  279. elif n_ellipsis == 0:
  280. # Note boolean masks must be the sole index, which we check for
  281. # later on.
  282. if not key_has_mask and n_single_axes < self.ndim:
  283. raise IndexError(
  284. f"{self.ndim=}, but the multi-axes index only specifies "
  285. f"{n_single_axes} dimensions. If this was intentional, "
  286. "add a trailing ellipsis (...) which expands into as many "
  287. "slices (:) as necessary - this is what np.ndarray arrays "
  288. "implicitly do, but such flat indexing behaviour is not "
  289. "specified in the Array API."
  290. )
  291. if n_ellipsis == 0:
  292. indexed_shape = self.shape
  293. else:
  294. ellipsis_start = None
  295. for pos, i in enumerate(nonexpanding_key):
  296. if not (isinstance(i, Array) or isinstance(i, np.ndarray)):
  297. if i == Ellipsis:
  298. ellipsis_start = pos
  299. break
  300. assert ellipsis_start is not None # sanity check
  301. ellipsis_end = self.ndim - (n_single_axes - ellipsis_start)
  302. indexed_shape = (
  303. self.shape[:ellipsis_start] + self.shape[ellipsis_end:]
  304. )
  305. for i, side in zip(single_axes, indexed_shape):
  306. if isinstance(i, slice):
  307. if side == 0:
  308. f_range = "0 (or None)"
  309. else:
  310. f_range = f"between -{side} and {side - 1} (or None)"
  311. if i.start is not None:
  312. try:
  313. start = operator.index(i.start)
  314. except TypeError:
  315. pass # handled by ndarray
  316. else:
  317. if not (-side <= start <= side):
  318. raise IndexError(
  319. f"Slice {i} contains {start=}, but should be "
  320. f"{f_range} for an axis of size {side} "
  321. "(out-of-bounds starts are not specified in "
  322. "the Array API)"
  323. )
  324. if i.stop is not None:
  325. try:
  326. stop = operator.index(i.stop)
  327. except TypeError:
  328. pass # handled by ndarray
  329. else:
  330. if not (-side <= stop <= side):
  331. raise IndexError(
  332. f"Slice {i} contains {stop=}, but should be "
  333. f"{f_range} for an axis of size {side} "
  334. "(out-of-bounds stops are not specified in "
  335. "the Array API)"
  336. )
  337. elif isinstance(i, Array):
  338. if i.dtype in _boolean_dtypes and len(_key) != 1:
  339. assert isinstance(key, tuple) # sanity check
  340. raise IndexError(
  341. f"Single-axes index {i} is a boolean array and "
  342. f"{len(key)=}, but masking is only specified in the "
  343. "Array API when the array is the sole index."
  344. )
  345. elif i.dtype in _integer_dtypes and i.ndim != 0:
  346. raise IndexError(
  347. f"Single-axes index {i} is a non-zero-dimensional "
  348. "integer array, but advanced integer indexing is not "
  349. "specified in the Array API."
  350. )
  351. elif isinstance(i, tuple):
  352. raise IndexError(
  353. f"Single-axes index {i} is a tuple, but nested tuple "
  354. "indices are not specified in the Array API."
  355. )
  356. # Everything below this line is required by the spec.
  357. def __abs__(self: Array, /) -> Array:
  358. """
  359. Performs the operation __abs__.
  360. """
  361. if self.dtype not in _numeric_dtypes:
  362. raise TypeError("Only numeric dtypes are allowed in __abs__")
  363. res = self._array.__abs__()
  364. return self.__class__._new(res)
  365. def __add__(self: Array, other: Union[int, float, Array], /) -> Array:
  366. """
  367. Performs the operation __add__.
  368. """
  369. other = self._check_allowed_dtypes(other, "numeric", "__add__")
  370. if other is NotImplemented:
  371. return other
  372. self, other = self._normalize_two_args(self, other)
  373. res = self._array.__add__(other._array)
  374. return self.__class__._new(res)
  375. def __and__(self: Array, other: Union[int, bool, Array], /) -> Array:
  376. """
  377. Performs the operation __and__.
  378. """
  379. other = self._check_allowed_dtypes(other, "integer or boolean", "__and__")
  380. if other is NotImplemented:
  381. return other
  382. self, other = self._normalize_two_args(self, other)
  383. res = self._array.__and__(other._array)
  384. return self.__class__._new(res)
  385. def __array_namespace__(
  386. self: Array, /, *, api_version: Optional[str] = None
  387. ) -> types.ModuleType:
  388. if api_version is not None and not api_version.startswith("2021."):
  389. raise ValueError(f"Unrecognized array API version: {api_version!r}")
  390. return array_api
  391. def __bool__(self: Array, /) -> bool:
  392. """
  393. Performs the operation __bool__.
  394. """
  395. # Note: This is an error here.
  396. if self._array.ndim != 0:
  397. raise TypeError("bool is only allowed on arrays with 0 dimensions")
  398. if self.dtype not in _boolean_dtypes:
  399. raise ValueError("bool is only allowed on boolean arrays")
  400. res = self._array.__bool__()
  401. return res
  402. def __dlpack__(self: Array, /, *, stream: None = None) -> PyCapsule:
  403. """
  404. Performs the operation __dlpack__.
  405. """
  406. return self._array.__dlpack__(stream=stream)
  407. def __dlpack_device__(self: Array, /) -> Tuple[IntEnum, int]:
  408. """
  409. Performs the operation __dlpack_device__.
  410. """
  411. # Note: device support is required for this
  412. return self._array.__dlpack_device__()
  413. def __eq__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
  414. """
  415. Performs the operation __eq__.
  416. """
  417. # Even though "all" dtypes are allowed, we still require them to be
  418. # promotable with each other.
  419. other = self._check_allowed_dtypes(other, "all", "__eq__")
  420. if other is NotImplemented:
  421. return other
  422. self, other = self._normalize_two_args(self, other)
  423. res = self._array.__eq__(other._array)
  424. return self.__class__._new(res)
  425. def __float__(self: Array, /) -> float:
  426. """
  427. Performs the operation __float__.
  428. """
  429. # Note: This is an error here.
  430. if self._array.ndim != 0:
  431. raise TypeError("float is only allowed on arrays with 0 dimensions")
  432. if self.dtype not in _floating_dtypes:
  433. raise ValueError("float is only allowed on floating-point arrays")
  434. res = self._array.__float__()
  435. return res
  436. def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
  437. """
  438. Performs the operation __floordiv__.
  439. """
  440. other = self._check_allowed_dtypes(other, "numeric", "__floordiv__")
  441. if other is NotImplemented:
  442. return other
  443. self, other = self._normalize_two_args(self, other)
  444. res = self._array.__floordiv__(other._array)
  445. return self.__class__._new(res)
  446. def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
  447. """
  448. Performs the operation __ge__.
  449. """
  450. other = self._check_allowed_dtypes(other, "numeric", "__ge__")
  451. if other is NotImplemented:
  452. return other
  453. self, other = self._normalize_two_args(self, other)
  454. res = self._array.__ge__(other._array)
  455. return self.__class__._new(res)
  456. def __getitem__(
  457. self: Array,
  458. key: Union[
  459. int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], Array
  460. ],
  461. /,
  462. ) -> Array:
  463. """
  464. Performs the operation __getitem__.
  465. """
  466. # Note: Only indices required by the spec are allowed. See the
  467. # docstring of _validate_index
  468. self._validate_index(key)
  469. if isinstance(key, Array):
  470. # Indexing self._array with array_api arrays can be erroneous
  471. key = key._array
  472. res = self._array.__getitem__(key)
  473. return self._new(res)
  474. def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
  475. """
  476. Performs the operation __gt__.
  477. """
  478. other = self._check_allowed_dtypes(other, "numeric", "__gt__")
  479. if other is NotImplemented:
  480. return other
  481. self, other = self._normalize_two_args(self, other)
  482. res = self._array.__gt__(other._array)
  483. return self.__class__._new(res)
  484. def __int__(self: Array, /) -> int:
  485. """
  486. Performs the operation __int__.
  487. """
  488. # Note: This is an error here.
  489. if self._array.ndim != 0:
  490. raise TypeError("int is only allowed on arrays with 0 dimensions")
  491. if self.dtype not in _integer_dtypes:
  492. raise ValueError("int is only allowed on integer arrays")
  493. res = self._array.__int__()
  494. return res
  495. def __index__(self: Array, /) -> int:
  496. """
  497. Performs the operation __index__.
  498. """
  499. res = self._array.__index__()
  500. return res
  501. def __invert__(self: Array, /) -> Array:
  502. """
  503. Performs the operation __invert__.
  504. """
  505. if self.dtype not in _integer_or_boolean_dtypes:
  506. raise TypeError("Only integer or boolean dtypes are allowed in __invert__")
  507. res = self._array.__invert__()
  508. return self.__class__._new(res)
  509. def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
  510. """
  511. Performs the operation __le__.
  512. """
  513. other = self._check_allowed_dtypes(other, "numeric", "__le__")
  514. if other is NotImplemented:
  515. return other
  516. self, other = self._normalize_two_args(self, other)
  517. res = self._array.__le__(other._array)
  518. return self.__class__._new(res)
  519. def __lshift__(self: Array, other: Union[int, Array], /) -> Array:
  520. """
  521. Performs the operation __lshift__.
  522. """
  523. other = self._check_allowed_dtypes(other, "integer", "__lshift__")
  524. if other is NotImplemented:
  525. return other
  526. self, other = self._normalize_two_args(self, other)
  527. res = self._array.__lshift__(other._array)
  528. return self.__class__._new(res)
  529. def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
  530. """
  531. Performs the operation __lt__.
  532. """
  533. other = self._check_allowed_dtypes(other, "numeric", "__lt__")
  534. if other is NotImplemented:
  535. return other
  536. self, other = self._normalize_two_args(self, other)
  537. res = self._array.__lt__(other._array)
  538. return self.__class__._new(res)
  539. def __matmul__(self: Array, other: Array, /) -> Array:
  540. """
  541. Performs the operation __matmul__.
  542. """
  543. # matmul is not defined for scalars, but without this, we may get
  544. # the wrong error message from asarray.
  545. other = self._check_allowed_dtypes(other, "numeric", "__matmul__")
  546. if other is NotImplemented:
  547. return other
  548. res = self._array.__matmul__(other._array)
  549. return self.__class__._new(res)
  550. def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
  551. """
  552. Performs the operation __mod__.
  553. """
  554. other = self._check_allowed_dtypes(other, "numeric", "__mod__")
  555. if other is NotImplemented:
  556. return other
  557. self, other = self._normalize_two_args(self, other)
  558. res = self._array.__mod__(other._array)
  559. return self.__class__._new(res)
  560. def __mul__(self: Array, other: Union[int, float, Array], /) -> Array:
  561. """
  562. Performs the operation __mul__.
  563. """
  564. other = self._check_allowed_dtypes(other, "numeric", "__mul__")
  565. if other is NotImplemented:
  566. return other
  567. self, other = self._normalize_two_args(self, other)
  568. res = self._array.__mul__(other._array)
  569. return self.__class__._new(res)
  570. def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
  571. """
  572. Performs the operation __ne__.
  573. """
  574. other = self._check_allowed_dtypes(other, "all", "__ne__")
  575. if other is NotImplemented:
  576. return other
  577. self, other = self._normalize_two_args(self, other)
  578. res = self._array.__ne__(other._array)
  579. return self.__class__._new(res)
  580. def __neg__(self: Array, /) -> Array:
  581. """
  582. Performs the operation __neg__.
  583. """
  584. if self.dtype not in _numeric_dtypes:
  585. raise TypeError("Only numeric dtypes are allowed in __neg__")
  586. res = self._array.__neg__()
  587. return self.__class__._new(res)
  588. def __or__(self: Array, other: Union[int, bool, Array], /) -> Array:
  589. """
  590. Performs the operation __or__.
  591. """
  592. other = self._check_allowed_dtypes(other, "integer or boolean", "__or__")
  593. if other is NotImplemented:
  594. return other
  595. self, other = self._normalize_two_args(self, other)
  596. res = self._array.__or__(other._array)
  597. return self.__class__._new(res)
  598. def __pos__(self: Array, /) -> Array:
  599. """
  600. Performs the operation __pos__.
  601. """
  602. if self.dtype not in _numeric_dtypes:
  603. raise TypeError("Only numeric dtypes are allowed in __pos__")
  604. res = self._array.__pos__()
  605. return self.__class__._new(res)
  606. def __pow__(self: Array, other: Union[int, float, Array], /) -> Array:
  607. """
  608. Performs the operation __pow__.
  609. """
  610. from ._elementwise_functions import pow
  611. other = self._check_allowed_dtypes(other, "numeric", "__pow__")
  612. if other is NotImplemented:
  613. return other
  614. # Note: NumPy's __pow__ does not follow type promotion rules for 0-d
  615. # arrays, so we use pow() here instead.
  616. return pow(self, other)
  617. def __rshift__(self: Array, other: Union[int, Array], /) -> Array:
  618. """
  619. Performs the operation __rshift__.
  620. """
  621. other = self._check_allowed_dtypes(other, "integer", "__rshift__")
  622. if other is NotImplemented:
  623. return other
  624. self, other = self._normalize_two_args(self, other)
  625. res = self._array.__rshift__(other._array)
  626. return self.__class__._new(res)
  627. def __setitem__(
  628. self,
  629. key: Union[
  630. int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], Array
  631. ],
  632. value: Union[int, float, bool, Array],
  633. /,
  634. ) -> None:
  635. """
  636. Performs the operation __setitem__.
  637. """
  638. # Note: Only indices required by the spec are allowed. See the
  639. # docstring of _validate_index
  640. self._validate_index(key)
  641. if isinstance(key, Array):
  642. # Indexing self._array with array_api arrays can be erroneous
  643. key = key._array
  644. self._array.__setitem__(key, asarray(value)._array)
  645. def __sub__(self: Array, other: Union[int, float, Array], /) -> Array:
  646. """
  647. Performs the operation __sub__.
  648. """
  649. other = self._check_allowed_dtypes(other, "numeric", "__sub__")
  650. if other is NotImplemented:
  651. return other
  652. self, other = self._normalize_two_args(self, other)
  653. res = self._array.__sub__(other._array)
  654. return self.__class__._new(res)
  655. # PEP 484 requires int to be a subtype of float, but __truediv__ should
  656. # not accept int.
  657. def __truediv__(self: Array, other: Union[float, Array], /) -> Array:
  658. """
  659. Performs the operation __truediv__.
  660. """
  661. other = self._check_allowed_dtypes(other, "floating-point", "__truediv__")
  662. if other is NotImplemented:
  663. return other
  664. self, other = self._normalize_two_args(self, other)
  665. res = self._array.__truediv__(other._array)
  666. return self.__class__._new(res)
  667. def __xor__(self: Array, other: Union[int, bool, Array], /) -> Array:
  668. """
  669. Performs the operation __xor__.
  670. """
  671. other = self._check_allowed_dtypes(other, "integer or boolean", "__xor__")
  672. if other is NotImplemented:
  673. return other
  674. self, other = self._normalize_two_args(self, other)
  675. res = self._array.__xor__(other._array)
  676. return self.__class__._new(res)
  677. def __iadd__(self: Array, other: Union[int, float, Array], /) -> Array:
  678. """
  679. Performs the operation __iadd__.
  680. """
  681. other = self._check_allowed_dtypes(other, "numeric", "__iadd__")
  682. if other is NotImplemented:
  683. return other
  684. self._array.__iadd__(other._array)
  685. return self
  686. def __radd__(self: Array, other: Union[int, float, Array], /) -> Array:
  687. """
  688. Performs the operation __radd__.
  689. """
  690. other = self._check_allowed_dtypes(other, "numeric", "__radd__")
  691. if other is NotImplemented:
  692. return other
  693. self, other = self._normalize_two_args(self, other)
  694. res = self._array.__radd__(other._array)
  695. return self.__class__._new(res)
  696. def __iand__(self: Array, other: Union[int, bool, Array], /) -> Array:
  697. """
  698. Performs the operation __iand__.
  699. """
  700. other = self._check_allowed_dtypes(other, "integer or boolean", "__iand__")
  701. if other is NotImplemented:
  702. return other
  703. self._array.__iand__(other._array)
  704. return self
  705. def __rand__(self: Array, other: Union[int, bool, Array], /) -> Array:
  706. """
  707. Performs the operation __rand__.
  708. """
  709. other = self._check_allowed_dtypes(other, "integer or boolean", "__rand__")
  710. if other is NotImplemented:
  711. return other
  712. self, other = self._normalize_two_args(self, other)
  713. res = self._array.__rand__(other._array)
  714. return self.__class__._new(res)
  715. def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
  716. """
  717. Performs the operation __ifloordiv__.
  718. """
  719. other = self._check_allowed_dtypes(other, "numeric", "__ifloordiv__")
  720. if other is NotImplemented:
  721. return other
  722. self._array.__ifloordiv__(other._array)
  723. return self
  724. def __rfloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
  725. """
  726. Performs the operation __rfloordiv__.
  727. """
  728. other = self._check_allowed_dtypes(other, "numeric", "__rfloordiv__")
  729. if other is NotImplemented:
  730. return other
  731. self, other = self._normalize_two_args(self, other)
  732. res = self._array.__rfloordiv__(other._array)
  733. return self.__class__._new(res)
  734. def __ilshift__(self: Array, other: Union[int, Array], /) -> Array:
  735. """
  736. Performs the operation __ilshift__.
  737. """
  738. other = self._check_allowed_dtypes(other, "integer", "__ilshift__")
  739. if other is NotImplemented:
  740. return other
  741. self._array.__ilshift__(other._array)
  742. return self
  743. def __rlshift__(self: Array, other: Union[int, Array], /) -> Array:
  744. """
  745. Performs the operation __rlshift__.
  746. """
  747. other = self._check_allowed_dtypes(other, "integer", "__rlshift__")
  748. if other is NotImplemented:
  749. return other
  750. self, other = self._normalize_two_args(self, other)
  751. res = self._array.__rlshift__(other._array)
  752. return self.__class__._new(res)
  753. def __imatmul__(self: Array, other: Array, /) -> Array:
  754. """
  755. Performs the operation __imatmul__.
  756. """
  757. # Note: NumPy does not implement __imatmul__.
  758. # matmul is not defined for scalars, but without this, we may get
  759. # the wrong error message from asarray.
  760. other = self._check_allowed_dtypes(other, "numeric", "__imatmul__")
  761. if other is NotImplemented:
  762. return other
  763. # __imatmul__ can only be allowed when it would not change the shape
  764. # of self.
  765. other_shape = other.shape
  766. if self.shape == () or other_shape == ():
  767. raise ValueError("@= requires at least one dimension")
  768. if len(other_shape) == 1 or other_shape[-1] != other_shape[-2]:
  769. raise ValueError("@= cannot change the shape of the input array")
  770. self._array[:] = self._array.__matmul__(other._array)
  771. return self
  772. def __rmatmul__(self: Array, other: Array, /) -> Array:
  773. """
  774. Performs the operation __rmatmul__.
  775. """
  776. # matmul is not defined for scalars, but without this, we may get
  777. # the wrong error message from asarray.
  778. other = self._check_allowed_dtypes(other, "numeric", "__rmatmul__")
  779. if other is NotImplemented:
  780. return other
  781. res = self._array.__rmatmul__(other._array)
  782. return self.__class__._new(res)
  783. def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
  784. """
  785. Performs the operation __imod__.
  786. """
  787. other = self._check_allowed_dtypes(other, "numeric", "__imod__")
  788. if other is NotImplemented:
  789. return other
  790. self._array.__imod__(other._array)
  791. return self
  792. def __rmod__(self: Array, other: Union[int, float, Array], /) -> Array:
  793. """
  794. Performs the operation __rmod__.
  795. """
  796. other = self._check_allowed_dtypes(other, "numeric", "__rmod__")
  797. if other is NotImplemented:
  798. return other
  799. self, other = self._normalize_two_args(self, other)
  800. res = self._array.__rmod__(other._array)
  801. return self.__class__._new(res)
  802. def __imul__(self: Array, other: Union[int, float, Array], /) -> Array:
  803. """
  804. Performs the operation __imul__.
  805. """
  806. other = self._check_allowed_dtypes(other, "numeric", "__imul__")
  807. if other is NotImplemented:
  808. return other
  809. self._array.__imul__(other._array)
  810. return self
  811. def __rmul__(self: Array, other: Union[int, float, Array], /) -> Array:
  812. """
  813. Performs the operation __rmul__.
  814. """
  815. other = self._check_allowed_dtypes(other, "numeric", "__rmul__")
  816. if other is NotImplemented:
  817. return other
  818. self, other = self._normalize_two_args(self, other)
  819. res = self._array.__rmul__(other._array)
  820. return self.__class__._new(res)
  821. def __ior__(self: Array, other: Union[int, bool, Array], /) -> Array:
  822. """
  823. Performs the operation __ior__.
  824. """
  825. other = self._check_allowed_dtypes(other, "integer or boolean", "__ior__")
  826. if other is NotImplemented:
  827. return other
  828. self._array.__ior__(other._array)
  829. return self
  830. def __ror__(self: Array, other: Union[int, bool, Array], /) -> Array:
  831. """
  832. Performs the operation __ror__.
  833. """
  834. other = self._check_allowed_dtypes(other, "integer or boolean", "__ror__")
  835. if other is NotImplemented:
  836. return other
  837. self, other = self._normalize_two_args(self, other)
  838. res = self._array.__ror__(other._array)
  839. return self.__class__._new(res)
  840. def __ipow__(self: Array, other: Union[int, float, Array], /) -> Array:
  841. """
  842. Performs the operation __ipow__.
  843. """
  844. other = self._check_allowed_dtypes(other, "numeric", "__ipow__")
  845. if other is NotImplemented:
  846. return other
  847. self._array.__ipow__(other._array)
  848. return self
  849. def __rpow__(self: Array, other: Union[int, float, Array], /) -> Array:
  850. """
  851. Performs the operation __rpow__.
  852. """
  853. from ._elementwise_functions import pow
  854. other = self._check_allowed_dtypes(other, "numeric", "__rpow__")
  855. if other is NotImplemented:
  856. return other
  857. # Note: NumPy's __pow__ does not follow the spec type promotion rules
  858. # for 0-d arrays, so we use pow() here instead.
  859. return pow(other, self)
  860. def __irshift__(self: Array, other: Union[int, Array], /) -> Array:
  861. """
  862. Performs the operation __irshift__.
  863. """
  864. other = self._check_allowed_dtypes(other, "integer", "__irshift__")
  865. if other is NotImplemented:
  866. return other
  867. self._array.__irshift__(other._array)
  868. return self
  869. def __rrshift__(self: Array, other: Union[int, Array], /) -> Array:
  870. """
  871. Performs the operation __rrshift__.
  872. """
  873. other = self._check_allowed_dtypes(other, "integer", "__rrshift__")
  874. if other is NotImplemented:
  875. return other
  876. self, other = self._normalize_two_args(self, other)
  877. res = self._array.__rrshift__(other._array)
  878. return self.__class__._new(res)
  879. def __isub__(self: Array, other: Union[int, float, Array], /) -> Array:
  880. """
  881. Performs the operation __isub__.
  882. """
  883. other = self._check_allowed_dtypes(other, "numeric", "__isub__")
  884. if other is NotImplemented:
  885. return other
  886. self._array.__isub__(other._array)
  887. return self
  888. def __rsub__(self: Array, other: Union[int, float, Array], /) -> Array:
  889. """
  890. Performs the operation __rsub__.
  891. """
  892. other = self._check_allowed_dtypes(other, "numeric", "__rsub__")
  893. if other is NotImplemented:
  894. return other
  895. self, other = self._normalize_two_args(self, other)
  896. res = self._array.__rsub__(other._array)
  897. return self.__class__._new(res)
  898. def __itruediv__(self: Array, other: Union[float, Array], /) -> Array:
  899. """
  900. Performs the operation __itruediv__.
  901. """
  902. other = self._check_allowed_dtypes(other, "floating-point", "__itruediv__")
  903. if other is NotImplemented:
  904. return other
  905. self._array.__itruediv__(other._array)
  906. return self
  907. def __rtruediv__(self: Array, other: Union[float, Array], /) -> Array:
  908. """
  909. Performs the operation __rtruediv__.
  910. """
  911. other = self._check_allowed_dtypes(other, "floating-point", "__rtruediv__")
  912. if other is NotImplemented:
  913. return other
  914. self, other = self._normalize_two_args(self, other)
  915. res = self._array.__rtruediv__(other._array)
  916. return self.__class__._new(res)
  917. def __ixor__(self: Array, other: Union[int, bool, Array], /) -> Array:
  918. """
  919. Performs the operation __ixor__.
  920. """
  921. other = self._check_allowed_dtypes(other, "integer or boolean", "__ixor__")
  922. if other is NotImplemented:
  923. return other
  924. self._array.__ixor__(other._array)
  925. return self
  926. def __rxor__(self: Array, other: Union[int, bool, Array], /) -> Array:
  927. """
  928. Performs the operation __rxor__.
  929. """
  930. other = self._check_allowed_dtypes(other, "integer or boolean", "__rxor__")
  931. if other is NotImplemented:
  932. return other
  933. self, other = self._normalize_two_args(self, other)
  934. res = self._array.__rxor__(other._array)
  935. return self.__class__._new(res)
  936. def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
  937. if stream is not None:
  938. raise ValueError("The stream argument to to_device() is not supported")
  939. if device == 'cpu':
  940. return self
  941. raise ValueError(f"Unsupported device {device!r}")
  942. @property
  943. def dtype(self) -> Dtype:
  944. """
  945. Array API compatible wrapper for :py:meth:`np.ndarray.dtype <numpy.ndarray.dtype>`.
  946. See its docstring for more information.
  947. """
  948. return self._array.dtype
  949. @property
  950. def device(self) -> Device:
  951. return "cpu"
  952. # Note: mT is new in array API spec (see matrix_transpose)
  953. @property
  954. def mT(self) -> Array:
  955. from .linalg import matrix_transpose
  956. return matrix_transpose(self)
  957. @property
  958. def ndim(self) -> int:
  959. """
  960. Array API compatible wrapper for :py:meth:`np.ndarray.ndim <numpy.ndarray.ndim>`.
  961. See its docstring for more information.
  962. """
  963. return self._array.ndim
  964. @property
  965. def shape(self) -> Tuple[int, ...]:
  966. """
  967. Array API compatible wrapper for :py:meth:`np.ndarray.shape <numpy.ndarray.shape>`.
  968. See its docstring for more information.
  969. """
  970. return self._array.shape
  971. @property
  972. def size(self) -> int:
  973. """
  974. Array API compatible wrapper for :py:meth:`np.ndarray.size <numpy.ndarray.size>`.
  975. See its docstring for more information.
  976. """
  977. return self._array.size
  978. @property
  979. def T(self) -> Array:
  980. """
  981. Array API compatible wrapper for :py:meth:`np.ndarray.T <numpy.ndarray.T>`.
  982. See its docstring for more information.
  983. """
  984. # Note: T only works on 2-dimensional arrays. See the corresponding
  985. # note in the specification:
  986. # https://data-apis.org/array-api/latest/API_specification/array_object.html#t
  987. if self.ndim != 2:
  988. raise ValueError("x.T requires x to have 2 dimensions. Use x.mT to transpose stacks of matrices and permute_dims() to permute dimensions.")
  989. return self.__class__._new(self._array.T)