array-api-strict 2.0.1__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.1 → array_api_strict-2.1}/PKG-INFO +1 -5
  2. {array_api_strict-2.0.1 → array_api_strict-2.1}/README.md +0 -4
  3. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/__init__.py +3 -0
  4. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_array_object.py +120 -50
  5. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_creation_functions.py +37 -21
  6. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_data_type_functions.py +5 -3
  7. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_dtypes.py +1 -1
  8. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_elementwise_functions.py +133 -67
  9. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_fft.py +17 -17
  10. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_flags.py +4 -8
  11. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_indexing_functions.py +3 -1
  12. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_info.py +2 -2
  13. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_linalg.py +31 -20
  14. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_linear_algebra_functions.py +13 -4
  15. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_manipulation_functions.py +28 -13
  16. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_searching_functions.py +13 -5
  17. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_set_functions.py +8 -7
  18. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_sorting_functions.py +2 -2
  19. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_statistical_functions.py +8 -8
  20. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_typing.py +2 -2
  21. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_utility_functions.py +2 -2
  22. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_version.py +3 -3
  23. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_array_object.py +16 -6
  24. array_api_strict-2.1/array_api_strict/tests/test_device_support.py +38 -0
  25. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_elementwise_functions.py +73 -7
  26. {array_api_strict-2.0.1 → 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.1 → array_api_strict-2.1}/array_api_strict/tests/test_linalg.py +14 -13
  29. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_sorting_functions.py +14 -0
  30. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_statistical_functions.py +15 -5
  31. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict.egg-info/PKG-INFO +1 -5
  32. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict.egg-info/SOURCES.txt +1 -0
  33. array_api_strict-2.0.1/array_api_strict/tests/test_indexing_functions.py +0 -24
  34. {array_api_strict-2.0.1 → array_api_strict-2.1}/LICENSE +0 -0
  35. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/_constants.py +0 -0
  36. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/__init__.py +0 -0
  37. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/conftest.py +0 -0
  38. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_creation_functions.py +0 -0
  39. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_data_type_functions.py +0 -0
  40. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_manipulation_functions.py +0 -0
  41. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_set_functions.py +0 -0
  42. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict/tests/test_validation.py +0 -0
  43. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict.egg-info/dependency_links.txt +0 -0
  44. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict.egg-info/requires.txt +0 -0
  45. {array_api_strict-2.0.1 → array_api_strict-2.1}/array_api_strict.egg-info/top_level.txt +0 -0
  46. {array_api_strict-2.0.1 → array_api_strict-2.1}/setup.cfg +0 -0
  47. {array_api_strict-2.0.1 → array_api_strict-2.1}/setup.py +0 -0
  48. {array_api_strict-2.0.1 → 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.1
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,7 +708,7 @@ 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
  """
@@ -686,85 +723,92 @@ class Array:
686
723
  # define __iter__, but it doesn't disallow it. The default Python
687
724
  # behavior is to implement iter as a[0], a[1], ... when __getitem__ is
688
725
  # implemented, which implies iteration on 1-D arrays.
689
- return (Array._new(i) for i in self._array)
726
+ return (Array._new(i, device=self.device) for i in self._array)
690
727
 
691
728
  def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
692
729
  """
693
730
  Performs the operation __le__.
694
731
  """
732
+ self._check_device(other)
695
733
  other = self._check_allowed_dtypes(other, "real numeric", "__le__")
696
734
  if other is NotImplemented:
697
735
  return other
698
736
  self, other = self._normalize_two_args(self, other)
699
737
  res = self._array.__le__(other._array)
700
- return self.__class__._new(res)
738
+ return self.__class__._new(res, device=self.device)
701
739
 
702
740
  def __lshift__(self: Array, other: Union[int, Array], /) -> Array:
703
741
  """
704
742
  Performs the operation __lshift__.
705
743
  """
744
+ self._check_device(other)
706
745
  other = self._check_allowed_dtypes(other, "integer", "__lshift__")
707
746
  if other is NotImplemented:
708
747
  return other
709
748
  self, other = self._normalize_two_args(self, other)
710
749
  res = self._array.__lshift__(other._array)
711
- return self.__class__._new(res)
750
+ return self.__class__._new(res, device=self.device)
712
751
 
713
752
  def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
714
753
  """
715
754
  Performs the operation __lt__.
716
755
  """
756
+ self._check_device(other)
717
757
  other = self._check_allowed_dtypes(other, "real numeric", "__lt__")
718
758
  if other is NotImplemented:
719
759
  return other
720
760
  self, other = self._normalize_two_args(self, other)
721
761
  res = self._array.__lt__(other._array)
722
- return self.__class__._new(res)
762
+ return self.__class__._new(res, device=self.device)
723
763
 
724
764
  def __matmul__(self: Array, other: Array, /) -> Array:
725
765
  """
726
766
  Performs the operation __matmul__.
727
767
  """
768
+ self._check_device(other)
728
769
  # matmul is not defined for scalars, but without this, we may get
729
770
  # the wrong error message from asarray.
730
771
  other = self._check_allowed_dtypes(other, "numeric", "__matmul__")
731
772
  if other is NotImplemented:
732
773
  return other
733
774
  res = self._array.__matmul__(other._array)
734
- return self.__class__._new(res)
775
+ return self.__class__._new(res, device=self.device)
735
776
 
736
777
  def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
737
778
  """
738
779
  Performs the operation __mod__.
739
780
  """
781
+ self._check_device(other)
740
782
  other = self._check_allowed_dtypes(other, "real numeric", "__mod__")
741
783
  if other is NotImplemented:
742
784
  return other
743
785
  self, other = self._normalize_two_args(self, other)
744
786
  res = self._array.__mod__(other._array)
745
- return self.__class__._new(res)
787
+ return self.__class__._new(res, device=self.device)
746
788
 
747
789
  def __mul__(self: Array, other: Union[int, float, Array], /) -> Array:
748
790
  """
749
791
  Performs the operation __mul__.
750
792
  """
793
+ self._check_device(other)
751
794
  other = self._check_allowed_dtypes(other, "numeric", "__mul__")
752
795
  if other is NotImplemented:
753
796
  return other
754
797
  self, other = self._normalize_two_args(self, other)
755
798
  res = self._array.__mul__(other._array)
756
- return self.__class__._new(res)
799
+ return self.__class__._new(res, device=self.device)
757
800
 
758
801
  def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
759
802
  """
760
803
  Performs the operation __ne__.
761
804
  """
805
+ self._check_device(other)
762
806
  other = self._check_allowed_dtypes(other, "all", "__ne__")
763
807
  if other is NotImplemented:
764
808
  return other
765
809
  self, other = self._normalize_two_args(self, other)
766
810
  res = self._array.__ne__(other._array)
767
- return self.__class__._new(res)
811
+ return self.__class__._new(res, device=self.device)
768
812
 
769
813
  def __neg__(self: Array, /) -> Array:
770
814
  """
@@ -773,18 +817,19 @@ class Array:
773
817
  if self.dtype not in _numeric_dtypes:
774
818
  raise TypeError("Only numeric dtypes are allowed in __neg__")
775
819
  res = self._array.__neg__()
776
- return self.__class__._new(res)
820
+ return self.__class__._new(res, device=self.device)
777
821
 
778
822
  def __or__(self: Array, other: Union[int, bool, Array], /) -> Array:
779
823
  """
780
824
  Performs the operation __or__.
781
825
  """
826
+ self._check_device(other)
782
827
  other = self._check_allowed_dtypes(other, "integer or boolean", "__or__")
783
828
  if other is NotImplemented:
784
829
  return other
785
830
  self, other = self._normalize_two_args(self, other)
786
831
  res = self._array.__or__(other._array)
787
- return self.__class__._new(res)
832
+ return self.__class__._new(res, device=self.device)
788
833
 
789
834
  def __pos__(self: Array, /) -> Array:
790
835
  """
@@ -793,7 +838,7 @@ class Array:
793
838
  if self.dtype not in _numeric_dtypes:
794
839
  raise TypeError("Only numeric dtypes are allowed in __pos__")
795
840
  res = self._array.__pos__()
796
- return self.__class__._new(res)
841
+ return self.__class__._new(res, device=self.device)
797
842
 
798
843
  def __pow__(self: Array, other: Union[int, float, Array], /) -> Array:
799
844
  """
@@ -801,6 +846,7 @@ class Array:
801
846
  """
802
847
  from ._elementwise_functions import pow
803
848
 
849
+ self._check_device(other)
804
850
  other = self._check_allowed_dtypes(other, "numeric", "__pow__")
805
851
  if other is NotImplemented:
806
852
  return other
@@ -812,12 +858,13 @@ class Array:
812
858
  """
813
859
  Performs the operation __rshift__.
814
860
  """
861
+ self._check_device(other)
815
862
  other = self._check_allowed_dtypes(other, "integer", "__rshift__")
816
863
  if other is NotImplemented:
817
864
  return other
818
865
  self, other = self._normalize_two_args(self, other)
819
866
  res = self._array.__rshift__(other._array)
820
- return self.__class__._new(res)
867
+ return self.__class__._new(res, device=self.device)
821
868
 
822
869
  def __setitem__(
823
870
  self,
@@ -842,12 +889,13 @@ class Array:
842
889
  """
843
890
  Performs the operation __sub__.
844
891
  """
892
+ self._check_device(other)
845
893
  other = self._check_allowed_dtypes(other, "numeric", "__sub__")
846
894
  if other is NotImplemented:
847
895
  return other
848
896
  self, other = self._normalize_two_args(self, other)
849
897
  res = self._array.__sub__(other._array)
850
- return self.__class__._new(res)
898
+ return self.__class__._new(res, device=self.device)
851
899
 
852
900
  # PEP 484 requires int to be a subtype of float, but __truediv__ should
853
901
  # not accept int.
@@ -855,28 +903,31 @@ class Array:
855
903
  """
856
904
  Performs the operation __truediv__.
857
905
  """
906
+ self._check_device(other)
858
907
  other = self._check_allowed_dtypes(other, "floating-point", "__truediv__")
859
908
  if other is NotImplemented:
860
909
  return other
861
910
  self, other = self._normalize_two_args(self, other)
862
911
  res = self._array.__truediv__(other._array)
863
- return self.__class__._new(res)
912
+ return self.__class__._new(res, device=self.device)
864
913
 
865
914
  def __xor__(self: Array, other: Union[int, bool, Array], /) -> Array:
866
915
  """
867
916
  Performs the operation __xor__.
868
917
  """
918
+ self._check_device(other)
869
919
  other = self._check_allowed_dtypes(other, "integer or boolean", "__xor__")
870
920
  if other is NotImplemented:
871
921
  return other
872
922
  self, other = self._normalize_two_args(self, other)
873
923
  res = self._array.__xor__(other._array)
874
- return self.__class__._new(res)
924
+ return self.__class__._new(res, device=self.device)
875
925
 
876
926
  def __iadd__(self: Array, other: Union[int, float, Array], /) -> Array:
877
927
  """
878
928
  Performs the operation __iadd__.
879
929
  """
930
+ self._check_device(other)
880
931
  other = self._check_allowed_dtypes(other, "numeric", "__iadd__")
881
932
  if other is NotImplemented:
882
933
  return other
@@ -887,17 +938,19 @@ class Array:
887
938
  """
888
939
  Performs the operation __radd__.
889
940
  """
941
+ self._check_device(other)
890
942
  other = self._check_allowed_dtypes(other, "numeric", "__radd__")
891
943
  if other is NotImplemented:
892
944
  return other
893
945
  self, other = self._normalize_two_args(self, other)
894
946
  res = self._array.__radd__(other._array)
895
- return self.__class__._new(res)
947
+ return self.__class__._new(res, device=self.device)
896
948
 
897
949
  def __iand__(self: Array, other: Union[int, bool, Array], /) -> Array:
898
950
  """
899
951
  Performs the operation __iand__.
900
952
  """
953
+ self._check_device(other)
901
954
  other = self._check_allowed_dtypes(other, "integer or boolean", "__iand__")
902
955
  if other is NotImplemented:
903
956
  return other
@@ -908,17 +961,19 @@ class Array:
908
961
  """
909
962
  Performs the operation __rand__.
910
963
  """
964
+ self._check_device(other)
911
965
  other = self._check_allowed_dtypes(other, "integer or boolean", "__rand__")
912
966
  if other is NotImplemented:
913
967
  return other
914
968
  self, other = self._normalize_two_args(self, other)
915
969
  res = self._array.__rand__(other._array)
916
- return self.__class__._new(res)
970
+ return self.__class__._new(res, device=self.device)
917
971
 
918
972
  def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
919
973
  """
920
974
  Performs the operation __ifloordiv__.
921
975
  """
976
+ self._check_device(other)
922
977
  other = self._check_allowed_dtypes(other, "real numeric", "__ifloordiv__")
923
978
  if other is NotImplemented:
924
979
  return other
@@ -929,17 +984,19 @@ class Array:
929
984
  """
930
985
  Performs the operation __rfloordiv__.
931
986
  """
987
+ self._check_device(other)
932
988
  other = self._check_allowed_dtypes(other, "real numeric", "__rfloordiv__")
933
989
  if other is NotImplemented:
934
990
  return other
935
991
  self, other = self._normalize_two_args(self, other)
936
992
  res = self._array.__rfloordiv__(other._array)
937
- return self.__class__._new(res)
993
+ return self.__class__._new(res, device=self.device)
938
994
 
939
995
  def __ilshift__(self: Array, other: Union[int, Array], /) -> Array:
940
996
  """
941
997
  Performs the operation __ilshift__.
942
998
  """
999
+ self._check_device(other)
943
1000
  other = self._check_allowed_dtypes(other, "integer", "__ilshift__")
944
1001
  if other is NotImplemented:
945
1002
  return other
@@ -950,12 +1007,13 @@ class Array:
950
1007
  """
951
1008
  Performs the operation __rlshift__.
952
1009
  """
1010
+ self._check_device(other)
953
1011
  other = self._check_allowed_dtypes(other, "integer", "__rlshift__")
954
1012
  if other is NotImplemented:
955
1013
  return other
956
1014
  self, other = self._normalize_two_args(self, other)
957
1015
  res = self._array.__rlshift__(other._array)
958
- return self.__class__._new(res)
1016
+ return self.__class__._new(res, device=self.device)
959
1017
 
960
1018
  def __imatmul__(self: Array, other: Array, /) -> Array:
961
1019
  """
@@ -966,8 +1024,9 @@ class Array:
966
1024
  other = self._check_allowed_dtypes(other, "numeric", "__imatmul__")
967
1025
  if other is NotImplemented:
968
1026
  return other
1027
+ self._check_device(other)
969
1028
  res = self._array.__imatmul__(other._array)
970
- return self.__class__._new(res)
1029
+ return self.__class__._new(res, device=self.device)
971
1030
 
972
1031
  def __rmatmul__(self: Array, other: Array, /) -> Array:
973
1032
  """
@@ -978,8 +1037,9 @@ class Array:
978
1037
  other = self._check_allowed_dtypes(other, "numeric", "__rmatmul__")
979
1038
  if other is NotImplemented:
980
1039
  return other
1040
+ self._check_device(other)
981
1041
  res = self._array.__rmatmul__(other._array)
982
- return self.__class__._new(res)
1042
+ return self.__class__._new(res, device=self.device)
983
1043
 
984
1044
  def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
985
1045
  """
@@ -998,9 +1058,10 @@ class Array:
998
1058
  other = self._check_allowed_dtypes(other, "real numeric", "__rmod__")
999
1059
  if other is NotImplemented:
1000
1060
  return other
1061
+ self._check_device(other)
1001
1062
  self, other = self._normalize_two_args(self, other)
1002
1063
  res = self._array.__rmod__(other._array)
1003
- return self.__class__._new(res)
1064
+ return self.__class__._new(res, device=self.device)
1004
1065
 
1005
1066
  def __imul__(self: Array, other: Union[int, float, Array], /) -> Array:
1006
1067
  """
@@ -1019,9 +1080,10 @@ class Array:
1019
1080
  other = self._check_allowed_dtypes(other, "numeric", "__rmul__")
1020
1081
  if other is NotImplemented:
1021
1082
  return other
1083
+ self._check_device(other)
1022
1084
  self, other = self._normalize_two_args(self, other)
1023
1085
  res = self._array.__rmul__(other._array)
1024
- return self.__class__._new(res)
1086
+ return self.__class__._new(res, device=self.device)
1025
1087
 
1026
1088
  def __ior__(self: Array, other: Union[int, bool, Array], /) -> Array:
1027
1089
  """
@@ -1037,12 +1099,13 @@ class Array:
1037
1099
  """
1038
1100
  Performs the operation __ror__.
1039
1101
  """
1102
+ self._check_device(other)
1040
1103
  other = self._check_allowed_dtypes(other, "integer or boolean", "__ror__")
1041
1104
  if other is NotImplemented:
1042
1105
  return other
1043
1106
  self, other = self._normalize_two_args(self, other)
1044
1107
  res = self._array.__ror__(other._array)
1045
- return self.__class__._new(res)
1108
+ return self.__class__._new(res, device=self.device)
1046
1109
 
1047
1110
  def __ipow__(self: Array, other: Union[int, float, Array], /) -> Array:
1048
1111
  """
@@ -1084,9 +1147,10 @@ class Array:
1084
1147
  other = self._check_allowed_dtypes(other, "integer", "__rrshift__")
1085
1148
  if other is NotImplemented:
1086
1149
  return other
1150
+ self._check_device(other)
1087
1151
  self, other = self._normalize_two_args(self, other)
1088
1152
  res = self._array.__rrshift__(other._array)
1089
- return self.__class__._new(res)
1153
+ return self.__class__._new(res, device=self.device)
1090
1154
 
1091
1155
  def __isub__(self: Array, other: Union[int, float, Array], /) -> Array:
1092
1156
  """
@@ -1105,9 +1169,10 @@ class Array:
1105
1169
  other = self._check_allowed_dtypes(other, "numeric", "__rsub__")
1106
1170
  if other is NotImplemented:
1107
1171
  return other
1172
+ self._check_device(other)
1108
1173
  self, other = self._normalize_two_args(self, other)
1109
1174
  res = self._array.__rsub__(other._array)
1110
- return self.__class__._new(res)
1175
+ return self.__class__._new(res, device=self.device)
1111
1176
 
1112
1177
  def __itruediv__(self: Array, other: Union[float, Array], /) -> Array:
1113
1178
  """
@@ -1126,9 +1191,10 @@ class Array:
1126
1191
  other = self._check_allowed_dtypes(other, "floating-point", "__rtruediv__")
1127
1192
  if other is NotImplemented:
1128
1193
  return other
1194
+ self._check_device(other)
1129
1195
  self, other = self._normalize_two_args(self, other)
1130
1196
  res = self._array.__rtruediv__(other._array)
1131
- return self.__class__._new(res)
1197
+ return self.__class__._new(res, device=self.device)
1132
1198
 
1133
1199
  def __ixor__(self: Array, other: Union[int, bool, Array], /) -> Array:
1134
1200
  """
@@ -1147,15 +1213,19 @@ class Array:
1147
1213
  other = self._check_allowed_dtypes(other, "integer or boolean", "__rxor__")
1148
1214
  if other is NotImplemented:
1149
1215
  return other
1216
+ self._check_device(other)
1150
1217
  self, other = self._normalize_two_args(self, other)
1151
1218
  res = self._array.__rxor__(other._array)
1152
- return self.__class__._new(res)
1219
+ return self.__class__._new(res, device=self.device)
1153
1220
 
1154
1221
  def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
1155
1222
  if stream is not None:
1156
1223
  raise ValueError("The stream argument to to_device() is not supported")
1157
- if device == CPU_DEVICE:
1224
+ if device == self._device:
1158
1225
  return self
1226
+ elif isinstance(device, Device):
1227
+ arr = np.asarray(self._array, copy=True)
1228
+ return self.__class__._new(arr, device=device)
1159
1229
  raise ValueError(f"Unsupported device {device!r}")
1160
1230
 
1161
1231
  @property
@@ -1169,7 +1239,7 @@ class Array:
1169
1239
 
1170
1240
  @property
1171
1241
  def device(self) -> Device:
1172
- return CPU_DEVICE
1242
+ return self._device
1173
1243
 
1174
1244
  # Note: mT is new in array API spec (see matrix_transpose)
1175
1245
  @property
@@ -1216,4 +1286,4 @@ class Array:
1216
1286
  # https://data-apis.org/array-api/latest/API_specification/array_object.html#t
1217
1287
  if self.ndim != 2:
1218
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.")
1219
- return self.__class__._new(self._array.T)
1289
+ return self.__class__._new(self._array.T, device=self.device)