Merge branch 'main' into compact-scalars

Author: crusaderky

array_api_strict/_array_object.py

@@ -234,14 +234,15 @@ def _check_allowed_dtypes(
         return other
 
     def _check_device(self, other: Array | complex) -> None:
-        """Check that other is on a device compatible with the current array"""
-        if isinstance(other, (bool, int, float, complex)):
-            return
-        elif isinstance(other, Array):
+        """Check that other is either a Python scalar or an array on a device
+        compatible with the current array.
+        """
+        if isinstance(other, Array):
             if self.device != other.device:
                 raise ValueError(f"Arrays from two different devices ({self.device} and {other.device}) can not be combined.")
-        else:
-            raise TypeError(f"Expected Array | python scalar; got {type(other)}")
+        # Disallow subclasses of Python scalars, such as np.float64 and np.complex128
+        elif type(other) not in (bool, int, float, complex):
+            raise TypeError(f"Expected Array or Python scalar; got {type(other)}")
 
     # Helper function to match the type promotion rules in the spec
     def _promote_scalar(self, scalar: complex) -> Array:
@@ -542,7 +543,7 @@ def __add__(self, other: Array | complex, /) -> Array:
         """
         Performs the operation __add__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__add__")
         if other is NotImplemented:
             return other
@@ -554,7 +555,7 @@ def __and__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __and__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer or boolean", "__and__")
         if other is NotImplemented:
             return other
@@ -651,7 +652,7 @@ def __eq__(self, other: Array | complex, /) -> Array:  # type: ignore[override]
         """
         Performs the operation __eq__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         # Even though "all" dtypes are allowed, we still require them to be
         # promotable with each other.
         other = self._check_allowed_dtypes(other, "all", "__eq__")
@@ -677,7 +678,7 @@ def __floordiv__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __floordiv__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__floordiv__")
         if other is NotImplemented:
             return other
@@ -689,7 +690,7 @@ def __ge__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __ge__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__ge__")
         if other is NotImplemented:
             return other
@@ -741,7 +742,7 @@ def __gt__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __gt__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__gt__")
         if other is NotImplemented:
             return other
@@ -796,7 +797,7 @@ def __le__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __le__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__le__")
         if other is NotImplemented:
             return other
@@ -808,7 +809,7 @@ def __lshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __lshift__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer", "__lshift__")
         if other is NotImplemented:
             return other
@@ -820,7 +821,7 @@ def __lt__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __lt__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__lt__")
         if other is NotImplemented:
             return other
@@ -832,7 +833,7 @@ def __matmul__(self, other: Array, /) -> Array:
         """
         Performs the operation __matmul__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         # matmul is not defined for scalars, but without this, we may get
         # the wrong error message from asarray.
         other = self._check_allowed_dtypes(other, "numeric", "__matmul__")
@@ -845,7 +846,7 @@ def __mod__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __mod__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__mod__")
         if other is NotImplemented:
             return other
@@ -857,7 +858,7 @@ def __mul__(self, other: Array | complex, /) -> Array:
         """
         Performs the operation __mul__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__mul__")
         if other is NotImplemented:
             return other
@@ -869,7 +870,7 @@ def __ne__(self, other: Array | complex, /) -> Array:  # type: ignore[override]
         """
         Performs the operation __ne__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "all", "__ne__")
         if other is NotImplemented:
             return other
@@ -890,7 +891,7 @@ def __or__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __or__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer or boolean", "__or__")
         if other is NotImplemented:
             return other
@@ -913,7 +914,7 @@ def __pow__(self, other: Array | complex, /) -> Array:
         """
         from ._elementwise_functions import pow  # type: ignore[attr-defined]
 
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__pow__")
         if other is NotImplemented:
             return other
@@ -925,7 +926,7 @@ def __rshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __rshift__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer", "__rshift__")
         if other is NotImplemented:
             return other
@@ -961,7 +962,7 @@ def __sub__(self, other: Array | complex, /) -> Array:
         """
         Performs the operation __sub__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__sub__")
         if other is NotImplemented:
             return other
@@ -975,7 +976,7 @@ def __truediv__(self, other: Array | complex, /) -> Array:
         """
         Performs the operation __truediv__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "floating-point", "__truediv__")
         if other is NotImplemented:
             return other
@@ -987,7 +988,7 @@ def __xor__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __xor__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer or boolean", "__xor__")
         if other is NotImplemented:
             return other
@@ -999,7 +1000,7 @@ def __iadd__(self, other: Array | complex, /) -> Array:
         """
         Performs the operation __iadd__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__iadd__")
         if other is NotImplemented:
             return other
@@ -1010,7 +1011,7 @@ def __radd__(self, other: Array | complex, /) -> Array:
         """
         Performs the operation __radd__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "numeric", "__radd__")
         if other is NotImplemented:
             return other
@@ -1022,7 +1023,7 @@ def __iand__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __iand__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer or boolean", "__iand__")
         if other is NotImplemented:
             return other
@@ -1033,7 +1034,7 @@ def __rand__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __rand__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer or boolean", "__rand__")
         if other is NotImplemented:
             return other
@@ -1045,7 +1046,7 @@ def __ifloordiv__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __ifloordiv__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__ifloordiv__")
         if other is NotImplemented:
             return other
@@ -1056,7 +1057,7 @@ def __rfloordiv__(self, other: Array | float, /) -> Array:
         """
         Performs the operation __rfloordiv__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "real numeric", "__rfloordiv__")
         if other is NotImplemented:
             return other
@@ -1068,7 +1069,7 @@ def __ilshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __ilshift__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer", "__ilshift__")
         if other is NotImplemented:
             return other
@@ -1079,7 +1080,7 @@ def __rlshift__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __rlshift__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer", "__rlshift__")
         if other is NotImplemented:
             return other
@@ -1096,7 +1097,7 @@ def __imatmul__(self, other: Array, /) -> Array:
         other = self._check_allowed_dtypes(other, "numeric", "__imatmul__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         res = self._array.__imatmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
@@ -1109,7 +1110,7 @@ def __rmatmul__(self, other: Array, /) -> Array:
         other = self._check_allowed_dtypes(other, "numeric", "__rmatmul__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         res = self._array.__rmatmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
@@ -1130,7 +1131,7 @@ def __rmod__(self, other: Array | float, /) -> Array:
         other = self._check_allowed_dtypes(other, "real numeric", "__rmod__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         self, other = self._normalize_two_args(self, other)
         res = self._array.__rmod__(other._array)
         return self.__class__._new(res, device=self.device)
@@ -1152,7 +1153,7 @@ def __rmul__(self, other: Array | complex, /) -> Array:
         other = self._check_allowed_dtypes(other, "numeric", "__rmul__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         self, other = self._normalize_two_args(self, other)
         res = self._array.__rmul__(other._array)
         return self.__class__._new(res, device=self.device)
@@ -1171,7 +1172,7 @@ def __ror__(self, other: Array | int, /) -> Array:
         """
         Performs the operation __ror__.
         """
-        self._check_device(other)
+        self._check_type_device(other)
         other = self._check_allowed_dtypes(other, "integer or boolean", "__ror__")
         if other is NotImplemented:
             return other
@@ -1219,7 +1220,7 @@ def __rrshift__(self, other: Array | int, /) -> Array:
         other = self._check_allowed_dtypes(other, "integer", "__rrshift__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         self, other = self._normalize_two_args(self, other)
         res = self._array.__rrshift__(other._array)
         return self.__class__._new(res, device=self.device)
@@ -1241,7 +1242,7 @@ def __rsub__(self, other: Array | complex, /) -> Array:
         other = self._check_allowed_dtypes(other, "numeric", "__rsub__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         self, other = self._normalize_two_args(self, other)
         res = self._array.__rsub__(other._array)
         return self.__class__._new(res, device=self.device)
@@ -1263,7 +1264,7 @@ def __rtruediv__(self, other: Array | complex, /) -> Array:
         other = self._check_allowed_dtypes(other, "floating-point", "__rtruediv__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         self, other = self._normalize_two_args(self, other)
         res = self._array.__rtruediv__(other._array)
         return self.__class__._new(res, device=self.device)
@@ -1285,7 +1286,7 @@ def __rxor__(self, other: Array | int, /) -> Array:
         other = self._check_allowed_dtypes(other, "integer or boolean", "__rxor__")
         if other is NotImplemented:
             return other
-        self._check_device(other)
+        self._check_type_device(other)
         self, other = self._normalize_two_args(self, other)
         res = self._array.__rxor__(other._array)
         return self.__class__._new(res, device=self.device)