array-api-strict 2.0__tar.gz → 2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {array_api_strict-2.0 → array_api_strict-2.1}/PKG-INFO +1 -5
  2. {array_api_strict-2.0 → array_api_strict-2.1}/README.md +0 -4
  3. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/__init__.py +3 -0
  4. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_array_object.py +129 -53
  5. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_creation_functions.py +37 -21
  6. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_data_type_functions.py +5 -3
  7. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_dtypes.py +1 -1
  8. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_elementwise_functions.py +133 -67
  9. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_fft.py +17 -17
  10. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_flags.py +19 -10
  11. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_indexing_functions.py +3 -1
  12. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_info.py +2 -2
  13. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_linalg.py +31 -20
  14. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_linear_algebra_functions.py +13 -4
  15. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_manipulation_functions.py +28 -13
  16. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_searching_functions.py +13 -5
  17. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_set_functions.py +8 -7
  18. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_sorting_functions.py +2 -2
  19. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_statistical_functions.py +8 -8
  20. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_typing.py +2 -2
  21. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_utility_functions.py +2 -2
  22. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_version.py +3 -3
  23. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_array_object.py +24 -8
  24. array_api_strict-2.1/array_api_strict/tests/test_device_support.py +38 -0
  25. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_elementwise_functions.py +73 -7
  26. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_flags.py +57 -48
  27. array_api_strict-2.1/array_api_strict/tests/test_indexing_functions.py +44 -0
  28. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_linalg.py +14 -13
  29. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_sorting_functions.py +14 -0
  30. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_statistical_functions.py +15 -5
  31. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict.egg-info/PKG-INFO +1 -5
  32. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict.egg-info/SOURCES.txt +1 -0
  33. array_api_strict-2.0/array_api_strict/tests/test_indexing_functions.py +0 -24
  34. {array_api_strict-2.0 → array_api_strict-2.1}/LICENSE +0 -0
  35. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/_constants.py +0 -0
  36. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/__init__.py +0 -0
  37. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/conftest.py +0 -0
  38. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_creation_functions.py +0 -0
  39. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_data_type_functions.py +0 -0
  40. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_manipulation_functions.py +0 -0
  41. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_set_functions.py +0 -0
  42. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict/tests/test_validation.py +0 -0
  43. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict.egg-info/dependency_links.txt +0 -0
  44. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict.egg-info/requires.txt +0 -0
  45. {array_api_strict-2.0 → array_api_strict-2.1}/array_api_strict.egg-info/top_level.txt +0 -0
  46. {array_api_strict-2.0 → array_api_strict-2.1}/setup.cfg +0 -0
  47. {array_api_strict-2.0 → array_api_strict-2.1}/setup.py +0 -0
  48. {array_api_strict-2.0 → array_api_strict-2.1}/versioneer.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: array_api_strict
3
- Version: 2.0
3
+ Version: 2.1
4
4
  Summary: A strict, minimal implementation of the Python array API standard.
5
5
  Home-page: https://data-apis.org/array-api-strict/
6
6
  Author: Consortium for Python Data API Standards
@@ -34,8 +34,4 @@ libraries. Consuming library code should use the
34
34
  support the array API. Rather, it is intended to be used in the test suites of
35
35
  consuming libraries to test their array API usage.
36
36
 
37
- array-api-strict currently supports the 2022.12 version of the standard.
38
- 2023.12 support is planned and is tracked by [this
39
- issue](https://github.com/data-apis/array-api-strict/issues/25).
40
-
41
37
  See the documentation for more details https://data-apis.org/array-api-strict/
@@ -15,8 +15,4 @@ libraries. Consuming library code should use the
15
15
  support the array API. Rather, it is intended to be used in the test suites of
16
16
  consuming libraries to test their array API usage.
17
17
 
18
- array-api-strict currently supports the 2022.12 version of the standard.
19
- 2023.12 support is planned and is tracked by [this
20
- issue](https://github.com/data-apis/array-api-strict/issues/25).
21
-
22
18
  See the documentation for more details https://data-apis.org/array-api-strict/
@@ -309,6 +309,9 @@ from ._utility_functions import all, any
309
309
 
310
310
  __all__ += ["all", "any"]
311
311
 
312
+ from ._array_object import Device
313
+ __all__ += ["Device"]
314
+
312
315
  # Helper functions that are not part of the standard
313
316
 
314
317
  from ._flags import (
@@ -43,13 +43,26 @@ if TYPE_CHECKING:
43
43
 
44
44
  import numpy as np
45
45
 
46
- # Placeholder object to represent the "cpu" device (the only device NumPy
47
- # supports).
48
- class _cpu_device:
46
+ class Device:
47
+ def __init__(self, device="CPU_DEVICE"):
48
+ if device not in ("CPU_DEVICE", "device1", "device2"):
49
+ raise ValueError(f"The device '{device}' is not a valid choice.")
50
+ self._device = device
51
+
49
52
  def __repr__(self):
50
- return "CPU_DEVICE"
53
+ return f"array_api_strict.Device('{self._device}')"
54
+
55
+ def __eq__(self, other):
56
+ if not isinstance(other, Device):
57
+ return False
58
+ return self._device == other._device
59
+
60
+ def __hash__(self):
61
+ return hash(("Device", self._device))
62
+
51
63
 
52
- CPU_DEVICE = _cpu_device()
64
+ CPU_DEVICE = Device()
65
+ ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"))
53
66
 
54
67
  _default = object()
55
68
 
@@ -73,7 +86,7 @@ class Array:
73
86
  # Use a custom constructor instead of __init__, as manually initializing
74
87
  # this class is not supported API.
75
88
  @classmethod
76
- def _new(cls, x, /):
89
+ def _new(cls, x, /, device):
77
90
  """
78
91
  This is a private method for initializing the array API Array
79
92
  object.
@@ -95,6 +108,9 @@ class Array:
95
108
  )
96
109
  obj._array = x
97
110
  obj._dtype = _dtype
111
+ if device is None:
112
+ device = CPU_DEVICE
113
+ obj._device = device
98
114
  return obj
99
115
 
100
116
  # Prevent Array() from working
@@ -116,7 +132,11 @@ class Array:
116
132
  """
117
133
  Performs the operation __repr__.
118
134
  """
119
- suffix = f", dtype={self.dtype})"
135
+ suffix = f", dtype={self.dtype}"
136
+ if self.device != CPU_DEVICE:
137
+ suffix += f", device={self.device})"
138
+ else:
139
+ suffix += ")"
120
140
  if 0 in self.shape:
121
141
  prefix = "empty("
122
142
  mid = str(self.shape)
@@ -134,6 +154,8 @@ class Array:
134
154
  will be present in other implementations.
135
155
 
136
156
  """
157
+ if self._device != CPU_DEVICE:
158
+ raise RuntimeError(f"Can not convert array on the '{self._device}' device to a Numpy array.")
137
159
  # copy keyword is new in 2.0.0; for older versions don't use it
138
160
  # retry without that keyword.
139
161
  if np.__version__[0] < '2':
@@ -193,6 +215,14 @@ class Array:
193
215
 
194
216
  return other
195
217
 
218
+ def _check_device(self, other):
219
+ """Check that other is on a device compatible with the current array"""
220
+ if isinstance(other, (int, complex, float, bool)):
221
+ return
222
+ elif isinstance(other, Array):
223
+ if self.device != other.device:
224
+ raise ValueError(f"Arrays from two different devices ({self.device} and {other.device}) can not be combined.")
225
+
196
226
  # Helper function to match the type promotion rules in the spec
197
227
  def _promote_scalar(self, scalar):
198
228
  """
@@ -244,7 +274,7 @@ class Array:
244
274
  # behavior for integers within the bounds of the integer dtype.
245
275
  # Outside of those bounds we use the default NumPy behavior (either
246
276
  # cast or raise OverflowError).
247
- return Array._new(np.array(scalar, dtype=self.dtype._np_dtype))
277
+ return Array._new(np.array(scalar, dtype=self.dtype._np_dtype), device=CPU_DEVICE)
248
278
 
249
279
  @staticmethod
250
280
  def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
@@ -276,9 +306,9 @@ class Array:
276
306
  # performant. broadcast_to(x1._array, x2.shape) is much slower. We
277
307
  # could also manually type promote x2, but that is more complicated
278
308
  # and about the same performance as this.
279
- x1 = Array._new(x1._array[None])
309
+ x1 = Array._new(x1._array[None], device=x1.device)
280
310
  elif x2.ndim == 0 and x1.ndim != 0:
281
- x2 = Array._new(x2._array[None])
311
+ x2 = Array._new(x2._array[None], device=x2.device)
282
312
  return (x1, x2)
283
313
 
284
314
  # Note: A large fraction of allowed indices are disallowed here (see the
@@ -462,29 +492,31 @@ class Array:
462
492
  if self.dtype not in _numeric_dtypes:
463
493
  raise TypeError("Only numeric dtypes are allowed in __abs__")
464
494
  res = self._array.__abs__()
465
- return self.__class__._new(res)
495
+ return self.__class__._new(res, device=self.device)
466
496
 
467
497
  def __add__(self: Array, other: Union[int, float, Array], /) -> Array:
468
498
  """
469
499
  Performs the operation __add__.
470
500
  """
501
+ self._check_device(other)
471
502
  other = self._check_allowed_dtypes(other, "numeric", "__add__")
472
503
  if other is NotImplemented:
473
504
  return other
474
505
  self, other = self._normalize_two_args(self, other)
475
506
  res = self._array.__add__(other._array)
476
- return self.__class__._new(res)
507
+ return self.__class__._new(res, device=self.device)
477
508
 
478
509
  def __and__(self: Array, other: Union[int, bool, Array], /) -> Array:
479
510
  """
480
511
  Performs the operation __and__.
481
512
  """
513
+ self._check_device(other)
482
514
  other = self._check_allowed_dtypes(other, "integer or boolean", "__and__")
483
515
  if other is NotImplemented:
484
516
  return other
485
517
  self, other = self._normalize_two_args(self, other)
486
518
  res = self._array.__and__(other._array)
487
- return self.__class__._new(res)
519
+ return self.__class__._new(res, device=self.device)
488
520
 
489
521
  def __array_namespace__(
490
522
  self: Array, /, *, api_version: Optional[str] = None
@@ -568,6 +600,7 @@ class Array:
568
600
  """
569
601
  Performs the operation __eq__.
570
602
  """
603
+ self._check_device(other)
571
604
  # Even though "all" dtypes are allowed, we still require them to be
572
605
  # promotable with each other.
573
606
  other = self._check_allowed_dtypes(other, "all", "__eq__")
@@ -575,7 +608,7 @@ class Array:
575
608
  return other
576
609
  self, other = self._normalize_two_args(self, other)
577
610
  res = self._array.__eq__(other._array)
578
- return self.__class__._new(res)
611
+ return self.__class__._new(res, device=self.device)
579
612
 
580
613
  def __float__(self: Array, /) -> float:
581
614
  """
@@ -593,23 +626,25 @@ class Array:
593
626
  """
594
627
  Performs the operation __floordiv__.
595
628
  """
629
+ self._check_device(other)
596
630
  other = self._check_allowed_dtypes(other, "real numeric", "__floordiv__")
597
631
  if other is NotImplemented:
598
632
  return other
599
633
  self, other = self._normalize_two_args(self, other)
600
634
  res = self._array.__floordiv__(other._array)
601
- return self.__class__._new(res)
635
+ return self.__class__._new(res, device=self.device)
602
636
 
603
637
  def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
604
638
  """
605
639
  Performs the operation __ge__.
606
640
  """
641
+ self._check_device(other)
607
642
  other = self._check_allowed_dtypes(other, "real numeric", "__ge__")
608
643
  if other is NotImplemented:
609
644
  return other
610
645
  self, other = self._normalize_two_args(self, other)
611
646
  res = self._array.__ge__(other._array)
612
- return self.__class__._new(res)
647
+ return self.__class__._new(res, device=self.device)
613
648
 
614
649
  def __getitem__(
615
650
  self: Array,
@@ -625,6 +660,7 @@ class Array:
625
660
  """
626
661
  Performs the operation __getitem__.
627
662
  """
663
+ # XXX Does key have to be on the same device? Is there an exception for CPU_DEVICE?
628
664
  # Note: Only indices required by the spec are allowed. See the
629
665
  # docstring of _validate_index
630
666
  self._validate_index(key)
@@ -632,18 +668,19 @@ class Array:
632
668
  # Indexing self._array with array_api_strict arrays can be erroneous
633
669
  key = key._array
634
670
  res = self._array.__getitem__(key)
635
- return self._new(res)
671
+ return self._new(res, device=self.device)
636
672
 
637
673
  def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
638
674
  """
639
675
  Performs the operation __gt__.
640
676
  """
677
+ self._check_device(other)
641
678
  other = self._check_allowed_dtypes(other, "real numeric", "__gt__")
642
679
  if other is NotImplemented:
643
680
  return other
644
681
  self, other = self._normalize_two_args(self, other)
645
682
  res = self._array.__gt__(other._array)
646
- return self.__class__._new(res)
683
+ return self.__class__._new(res, device=other.device)
647
684
 
648
685
  def __int__(self: Array, /) -> int:
649
686
  """
@@ -671,94 +708,107 @@ class Array:
671
708
  if self.dtype not in _integer_or_boolean_dtypes:
672
709
  raise TypeError("Only integer or boolean dtypes are allowed in __invert__")
673
710
  res = self._array.__invert__()
674
- return self.__class__._new(res)
711
+ return self.__class__._new(res, device=self.device)
675
712
 
676
713
  def __iter__(self: Array, /):
677
714
  """
678
715
  Performs the operation __iter__.
679
716
  """
680
- # Manually disable iteration, since __getitem__ raises IndexError on
681
- # things like ones((3, 3))[0], which causes list(ones((3, 3))) to give
682
- # [].
683
- raise TypeError("array iteration is not allowed in array-api-strict")
717
+ # Manually disable iteration on higher dimensional arrays, since
718
+ # __getitem__ raises IndexError on things like ones((3, 3))[0], which
719
+ # causes list(ones((3, 3))) to give [].
720
+ if self.ndim > 1:
721
+ raise TypeError("array iteration is not allowed in array-api-strict")
722
+ # Allow iteration for 1-D arrays. The array API doesn't strictly
723
+ # define __iter__, but it doesn't disallow it. The default Python
724
+ # behavior is to implement iter as a[0], a[1], ... when __getitem__ is
725
+ # implemented, which implies iteration on 1-D arrays.
726
+ return (Array._new(i, device=self.device) for i in self._array)
684
727
 
685
728
  def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
686
729
  """
687
730
  Performs the operation __le__.
688
731
  """
732
+ self._check_device(other)
689
733
  other = self._check_allowed_dtypes(other, "real numeric", "__le__")
690
734
  if other is NotImplemented:
691
735
  return other
692
736
  self, other = self._normalize_two_args(self, other)
693
737
  res = self._array.__le__(other._array)
694
- return self.__class__._new(res)
738
+ return self.__class__._new(res, device=self.device)
695
739
 
696
740
  def __lshift__(self: Array, other: Union[int, Array], /) -> Array:
697
741
  """
698
742
  Performs the operation __lshift__.
699
743
  """
744
+ self._check_device(other)
700
745
  other = self._check_allowed_dtypes(other, "integer", "__lshift__")
701
746
  if other is NotImplemented:
702
747
  return other
703
748
  self, other = self._normalize_two_args(self, other)
704
749
  res = self._array.__lshift__(other._array)
705
- return self.__class__._new(res)
750
+ return self.__class__._new(res, device=self.device)
706
751
 
707
752
  def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
708
753
  """
709
754
  Performs the operation __lt__.
710
755
  """
756
+ self._check_device(other)
711
757
  other = self._check_allowed_dtypes(other, "real numeric", "__lt__")
712
758
  if other is NotImplemented:
713
759
  return other
714
760
  self, other = self._normalize_two_args(self, other)
715
761
  res = self._array.__lt__(other._array)
716
- return self.__class__._new(res)
762
+ return self.__class__._new(res, device=self.device)
717
763
 
718
764
  def __matmul__(self: Array, other: Array, /) -> Array:
719
765
  """
720
766
  Performs the operation __matmul__.
721
767
  """
768
+ self._check_device(other)
722
769
  # matmul is not defined for scalars, but without this, we may get
723
770
  # the wrong error message from asarray.
724
771
  other = self._check_allowed_dtypes(other, "numeric", "__matmul__")
725
772
  if other is NotImplemented:
726
773
  return other
727
774
  res = self._array.__matmul__(other._array)
728
- return self.__class__._new(res)
775
+ return self.__class__._new(res, device=self.device)
729
776
 
730
777
  def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
731
778
  """
732
779
  Performs the operation __mod__.
733
780
  """
781
+ self._check_device(other)
734
782
  other = self._check_allowed_dtypes(other, "real numeric", "__mod__")
735
783
  if other is NotImplemented:
736
784
  return other
737
785
  self, other = self._normalize_two_args(self, other)
738
786
  res = self._array.__mod__(other._array)
739
- return self.__class__._new(res)
787
+ return self.__class__._new(res, device=self.device)
740
788
 
741
789
  def __mul__(self: Array, other: Union[int, float, Array], /) -> Array:
742
790
  """
743
791
  Performs the operation __mul__.
744
792
  """
793
+ self._check_device(other)
745
794
  other = self._check_allowed_dtypes(other, "numeric", "__mul__")
746
795
  if other is NotImplemented:
747
796
  return other
748
797
  self, other = self._normalize_two_args(self, other)
749
798
  res = self._array.__mul__(other._array)
750
- return self.__class__._new(res)
799
+ return self.__class__._new(res, device=self.device)
751
800
 
752
801
  def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
753
802
  """
754
803
  Performs the operation __ne__.
755
804
  """
805
+ self._check_device(other)
756
806
  other = self._check_allowed_dtypes(other, "all", "__ne__")
757
807
  if other is NotImplemented:
758
808
  return other
759
809
  self, other = self._normalize_two_args(self, other)
760
810
  res = self._array.__ne__(other._array)
761
- return self.__class__._new(res)
811
+ return self.__class__._new(res, device=self.device)
762
812
 
763
813
  def __neg__(self: Array, /) -> Array:
764
814
  """
@@ -767,18 +817,19 @@ class Array:
767
817
  if self.dtype not in _numeric_dtypes:
768
818
  raise TypeError("Only numeric dtypes are allowed in __neg__")
769
819
  res = self._array.__neg__()
770
- return self.__class__._new(res)
820
+ return self.__class__._new(res, device=self.device)
771
821
 
772
822
  def __or__(self: Array, other: Union[int, bool, Array], /) -> Array:
773
823
  """
774
824
  Performs the operation __or__.
775
825
  """
826
+ self._check_device(other)
776
827
  other = self._check_allowed_dtypes(other, "integer or boolean", "__or__")
777
828
  if other is NotImplemented:
778
829
  return other
779
830
  self, other = self._normalize_two_args(self, other)
780
831
  res = self._array.__or__(other._array)
781
- return self.__class__._new(res)
832
+ return self.__class__._new(res, device=self.device)
782
833
 
783
834
  def __pos__(self: Array, /) -> Array:
784
835
  """
@@ -787,7 +838,7 @@ class Array:
787
838
  if self.dtype not in _numeric_dtypes:
788
839
  raise TypeError("Only numeric dtypes are allowed in __pos__")
789
840
  res = self._array.__pos__()
790
- return self.__class__._new(res)
841
+ return self.__class__._new(res, device=self.device)
791
842
 
792
843
  def __pow__(self: Array, other: Union[int, float, Array], /) -> Array:
793
844
  """
@@ -795,6 +846,7 @@ class Array:
795
846
  """
796
847
  from ._elementwise_functions import pow
797
848
 
849
+ self._check_device(other)
798
850
  other = self._check_allowed_dtypes(other, "numeric", "__pow__")
799
851
  if other is NotImplemented:
800
852
  return other
@@ -806,12 +858,13 @@ class Array:
806
858
  """
807
859
  Performs the operation __rshift__.
808
860
  """
861
+ self._check_device(other)
809
862
  other = self._check_allowed_dtypes(other, "integer", "__rshift__")
810
863
  if other is NotImplemented:
811
864
  return other
812
865
  self, other = self._normalize_two_args(self, other)
813
866
  res = self._array.__rshift__(other._array)
814
- return self.__class__._new(res)
867
+ return self.__class__._new(res, device=self.device)
815
868
 
816
869
  def __setitem__(
817
870
  self,
@@ -836,12 +889,13 @@ class Array:
836
889
  """
837
890
  Performs the operation __sub__.
838
891
  """
892
+ self._check_device(other)
839
893
  other = self._check_allowed_dtypes(other, "numeric", "__sub__")
840
894
  if other is NotImplemented:
841
895
  return other
842
896
  self, other = self._normalize_two_args(self, other)
843
897
  res = self._array.__sub__(other._array)
844
- return self.__class__._new(res)
898
+ return self.__class__._new(res, device=self.device)
845
899
 
846
900
  # PEP 484 requires int to be a subtype of float, but __truediv__ should
847
901
  # not accept int.
@@ -849,28 +903,31 @@ class Array:
849
903
  """
850
904
  Performs the operation __truediv__.
851
905
  """
906
+ self._check_device(other)
852
907
  other = self._check_allowed_dtypes(other, "floating-point", "__truediv__")
853
908
  if other is NotImplemented:
854
909
  return other
855
910
  self, other = self._normalize_two_args(self, other)
856
911
  res = self._array.__truediv__(other._array)
857
- return self.__class__._new(res)
912
+ return self.__class__._new(res, device=self.device)
858
913
 
859
914
  def __xor__(self: Array, other: Union[int, bool, Array], /) -> Array:
860
915
  """
861
916
  Performs the operation __xor__.
862
917
  """
918
+ self._check_device(other)
863
919
  other = self._check_allowed_dtypes(other, "integer or boolean", "__xor__")
864
920
  if other is NotImplemented:
865
921
  return other
866
922
  self, other = self._normalize_two_args(self, other)
867
923
  res = self._array.__xor__(other._array)
868
- return self.__class__._new(res)
924
+ return self.__class__._new(res, device=self.device)
869
925
 
870
926
  def __iadd__(self: Array, other: Union[int, float, Array], /) -> Array:
871
927
  """
872
928
  Performs the operation __iadd__.
873
929
  """
930
+ self._check_device(other)
874
931
  other = self._check_allowed_dtypes(other, "numeric", "__iadd__")
875
932
  if other is NotImplemented:
876
933
  return other
@@ -881,17 +938,19 @@ class Array:
881
938
  """
882
939
  Performs the operation __radd__.
883
940
  """
941
+ self._check_device(other)
884
942
  other = self._check_allowed_dtypes(other, "numeric", "__radd__")
885
943
  if other is NotImplemented:
886
944
  return other
887
945
  self, other = self._normalize_two_args(self, other)
888
946
  res = self._array.__radd__(other._array)
889
- return self.__class__._new(res)
947
+ return self.__class__._new(res, device=self.device)
890
948
 
891
949
  def __iand__(self: Array, other: Union[int, bool, Array], /) -> Array:
892
950
  """
893
951
  Performs the operation __iand__.
894
952
  """
953
+ self._check_device(other)
895
954
  other = self._check_allowed_dtypes(other, "integer or boolean", "__iand__")
896
955
  if other is NotImplemented:
897
956
  return other
@@ -902,17 +961,19 @@ class Array:
902
961
  """
903
962
  Performs the operation __rand__.
904
963
  """
964
+ self._check_device(other)
905
965
  other = self._check_allowed_dtypes(other, "integer or boolean", "__rand__")
906
966
  if other is NotImplemented:
907
967
  return other
908
968
  self, other = self._normalize_two_args(self, other)
909
969
  res = self._array.__rand__(other._array)
910
- return self.__class__._new(res)
970
+ return self.__class__._new(res, device=self.device)
911
971
 
912
972
  def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
913
973
  """
914
974
  Performs the operation __ifloordiv__.
915
975
  """
976
+ self._check_device(other)
916
977
  other = self._check_allowed_dtypes(other, "real numeric", "__ifloordiv__")
917
978
  if other is NotImplemented:
918
979
  return other
@@ -923,17 +984,19 @@ class Array:
923
984
  """
924
985
  Performs the operation __rfloordiv__.
925
986
  """
987
+ self._check_device(other)
926
988
  other = self._check_allowed_dtypes(other, "real numeric", "__rfloordiv__")
927
989
  if other is NotImplemented:
928
990
  return other
929
991
  self, other = self._normalize_two_args(self, other)
930
992
  res = self._array.__rfloordiv__(other._array)
931
- return self.__class__._new(res)
993
+ return self.__class__._new(res, device=self.device)
932
994
 
933
995
  def __ilshift__(self: Array, other: Union[int, Array], /) -> Array:
934
996
  """
935
997
  Performs the operation __ilshift__.
936
998
  """
999
+ self._check_device(other)
937
1000
  other = self._check_allowed_dtypes(other, "integer", "__ilshift__")
938
1001
  if other is NotImplemented:
939
1002
  return other
@@ -944,12 +1007,13 @@ class Array:
944
1007
  """
945
1008
  Performs the operation __rlshift__.
946
1009
  """
1010
+ self._check_device(other)
947
1011
  other = self._check_allowed_dtypes(other, "integer", "__rlshift__")
948
1012
  if other is NotImplemented:
949
1013
  return other
950
1014
  self, other = self._normalize_two_args(self, other)
951
1015
  res = self._array.__rlshift__(other._array)
952
- return self.__class__._new(res)
1016
+ return self.__class__._new(res, device=self.device)
953
1017
 
954
1018
  def __imatmul__(self: Array, other: Array, /) -> Array:
955
1019
  """
@@ -960,8 +1024,9 @@ class Array:
960
1024
  other = self._check_allowed_dtypes(other, "numeric", "__imatmul__")
961
1025
  if other is NotImplemented:
962
1026
  return other
1027
+ self._check_device(other)
963
1028
  res = self._array.__imatmul__(other._array)
964
- return self.__class__._new(res)
1029
+ return self.__class__._new(res, device=self.device)
965
1030
 
966
1031
  def __rmatmul__(self: Array, other: Array, /) -> Array:
967
1032
  """
@@ -972,8 +1037,9 @@ class Array:
972
1037
  other = self._check_allowed_dtypes(other, "numeric", "__rmatmul__")
973
1038
  if other is NotImplemented:
974
1039
  return other
1040
+ self._check_device(other)
975
1041
  res = self._array.__rmatmul__(other._array)
976
- return self.__class__._new(res)
1042
+ return self.__class__._new(res, device=self.device)
977
1043
 
978
1044
  def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
979
1045
  """
@@ -992,9 +1058,10 @@ class Array:
992
1058
  other = self._check_allowed_dtypes(other, "real numeric", "__rmod__")
993
1059
  if other is NotImplemented:
994
1060
  return other
1061
+ self._check_device(other)
995
1062
  self, other = self._normalize_two_args(self, other)
996
1063
  res = self._array.__rmod__(other._array)
997
- return self.__class__._new(res)
1064
+ return self.__class__._new(res, device=self.device)
998
1065
 
999
1066
  def __imul__(self: Array, other: Union[int, float, Array], /) -> Array:
1000
1067
  """
@@ -1013,9 +1080,10 @@ class Array:
1013
1080
  other = self._check_allowed_dtypes(other, "numeric", "__rmul__")
1014
1081
  if other is NotImplemented:
1015
1082
  return other
1083
+ self._check_device(other)
1016
1084
  self, other = self._normalize_two_args(self, other)
1017
1085
  res = self._array.__rmul__(other._array)
1018
- return self.__class__._new(res)
1086
+ return self.__class__._new(res, device=self.device)
1019
1087
 
1020
1088
  def __ior__(self: Array, other: Union[int, bool, Array], /) -> Array:
1021
1089
  """
@@ -1031,12 +1099,13 @@ class Array:
1031
1099
  """
1032
1100
  Performs the operation __ror__.
1033
1101
  """
1102
+ self._check_device(other)
1034
1103
  other = self._check_allowed_dtypes(other, "integer or boolean", "__ror__")
1035
1104
  if other is NotImplemented:
1036
1105
  return other
1037
1106
  self, other = self._normalize_two_args(self, other)
1038
1107
  res = self._array.__ror__(other._array)
1039
- return self.__class__._new(res)
1108
+ return self.__class__._new(res, device=self.device)
1040
1109
 
1041
1110
  def __ipow__(self: Array, other: Union[int, float, Array], /) -> Array:
1042
1111
  """
@@ -1078,9 +1147,10 @@ class Array:
1078
1147
  other = self._check_allowed_dtypes(other, "integer", "__rrshift__")
1079
1148
  if other is NotImplemented:
1080
1149
  return other
1150
+ self._check_device(other)
1081
1151
  self, other = self._normalize_two_args(self, other)
1082
1152
  res = self._array.__rrshift__(other._array)
1083
- return self.__class__._new(res)
1153
+ return self.__class__._new(res, device=self.device)
1084
1154
 
1085
1155
  def __isub__(self: Array, other: Union[int, float, Array], /) -> Array:
1086
1156
  """
@@ -1099,9 +1169,10 @@ class Array:
1099
1169
  other = self._check_allowed_dtypes(other, "numeric", "__rsub__")
1100
1170
  if other is NotImplemented:
1101
1171
  return other
1172
+ self._check_device(other)
1102
1173
  self, other = self._normalize_two_args(self, other)
1103
1174
  res = self._array.__rsub__(other._array)
1104
- return self.__class__._new(res)
1175
+ return self.__class__._new(res, device=self.device)
1105
1176
 
1106
1177
  def __itruediv__(self: Array, other: Union[float, Array], /) -> Array:
1107
1178
  """
@@ -1120,9 +1191,10 @@ class Array:
1120
1191
  other = self._check_allowed_dtypes(other, "floating-point", "__rtruediv__")
1121
1192
  if other is NotImplemented:
1122
1193
  return other
1194
+ self._check_device(other)
1123
1195
  self, other = self._normalize_two_args(self, other)
1124
1196
  res = self._array.__rtruediv__(other._array)
1125
- return self.__class__._new(res)
1197
+ return self.__class__._new(res, device=self.device)
1126
1198
 
1127
1199
  def __ixor__(self: Array, other: Union[int, bool, Array], /) -> Array:
1128
1200
  """
@@ -1141,15 +1213,19 @@ class Array:
1141
1213
  other = self._check_allowed_dtypes(other, "integer or boolean", "__rxor__")
1142
1214
  if other is NotImplemented:
1143
1215
  return other
1216
+ self._check_device(other)
1144
1217
  self, other = self._normalize_two_args(self, other)
1145
1218
  res = self._array.__rxor__(other._array)
1146
- return self.__class__._new(res)
1219
+ return self.__class__._new(res, device=self.device)
1147
1220
 
1148
1221
  def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
1149
1222
  if stream is not None:
1150
1223
  raise ValueError("The stream argument to to_device() is not supported")
1151
- if device == CPU_DEVICE:
1224
+ if device == self._device:
1152
1225
  return self
1226
+ elif isinstance(device, Device):
1227
+ arr = np.asarray(self._array, copy=True)
1228
+ return self.__class__._new(arr, device=device)
1153
1229
  raise ValueError(f"Unsupported device {device!r}")
1154
1230
 
1155
1231
  @property
@@ -1163,7 +1239,7 @@ class Array:
1163
1239
 
1164
1240
  @property
1165
1241
  def device(self) -> Device:
1166
- return CPU_DEVICE
1242
+ return self._device
1167
1243
 
1168
1244
  # Note: mT is new in array API spec (see matrix_transpose)
1169
1245
  @property
@@ -1210,4 +1286,4 @@ class Array:
1210
1286
  # https://data-apis.org/array-api/latest/API_specification/array_object.html#t
1211
1287
  if self.ndim != 2:
1212
1288
  raise ValueError("x.T requires x to have 2 dimensions. Use x.mT to transpose stacks of matrices and permute_dims() to permute dimensions.")
1213
- return self.__class__._new(self._array.T)
1289
+ return self.__class__._new(self._array.T, device=self.device)