Merge pull request #25370 from asmeurer/array_api-portability

ENH: Make `numpy.array_api` more portable

Original NumPy Commit: ad360324dbb3f43f0b40384e97e2f251b4140c2f

Author: rgommers

array_api_strict/_array_object.py

@@ -19,6 +19,7 @@
 from enum import IntEnum
 from ._creation_functions import asarray
 from ._dtypes import (
+    _DType,
     _all_dtypes,
     _boolean_dtypes,
     _integer_dtypes,
@@ -39,6 +40,13 @@
 
 import numpy as np
 
+# Placeholder object to represent the "cpu" device (the only device NumPy
+# supports).
+class _cpu_device:
+    def __repr__(self):
+        return "CPU_DEVICE"
+
+CPU_DEVICE = _cpu_device()
 
 class Array:
     """
@@ -75,11 +83,13 @@ def _new(cls, x, /):
         if isinstance(x, np.generic):
             # Convert the array scalar to a 0-D array
             x = np.asarray(x)
-        if x.dtype not in _all_dtypes:
+        _dtype = _DType(x.dtype)
+        if _dtype not in _all_dtypes:
             raise TypeError(
                 f"The array_api namespace does not support the dtype '{x.dtype}'"
             )
         obj._array = x
+        obj._dtype = _dtype
         return obj
 
     # Prevent Array() from working
@@ -101,7 +111,7 @@ def __repr__(self: Array, /) -> str:
         """
         Performs the operation __repr__.
         """
-        suffix = f", dtype={self.dtype.name})"
+        suffix = f", dtype={self.dtype})"
         if 0 in self.shape:
             prefix = "empty("
             mid = str(self.shape)
@@ -176,6 +186,8 @@ def _promote_scalar(self, scalar):
         integer that is too large to fit in a NumPy integer dtype, or
         TypeError when the scalar type is incompatible with the dtype of self.
         """
+        from ._data_type_functions import iinfo
+
         # Note: Only Python scalar types that match the array dtype are
         # allowed.
         if isinstance(scalar, bool):
@@ -189,7 +201,7 @@ def _promote_scalar(self, scalar):
                     "Python int scalars cannot be promoted with bool arrays"
                 )
             if self.dtype in _integer_dtypes:
-                info = np.iinfo(self.dtype)
+                info = iinfo(self.dtype)
                 if not (info.min <= scalar <= info.max):
                     raise OverflowError(
                         "Python int scalars must be within the bounds of the dtype for integer arrays"
@@ -215,7 +227,7 @@ def _promote_scalar(self, scalar):
         # behavior for integers within the bounds of the integer dtype.
         # Outside of those bounds we use the default NumPy behavior (either
         # cast or raise OverflowError).
-        return Array._new(np.array(scalar, self.dtype))
+        return Array._new(np.array(scalar, dtype=self.dtype._np_dtype))
 
     @staticmethod
     def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
@@ -325,7 +337,9 @@ def _validate_index(self, key):
         for i in _key:
             if i is not None:
                 nonexpanding_key.append(i)
-                if isinstance(i, Array) or isinstance(i, np.ndarray):
+                if isinstance(i, np.ndarray):
+                    raise IndexError("Index arrays for np.array_api must be np.array_api arrays")
+                if isinstance(i, Array):
                     if i.dtype in _boolean_dtypes:
                         key_has_mask = True
                     single_axes.append(i)
@@ -1067,7 +1081,7 @@ def __rxor__(self: Array, other: Union[int, bool, Array], /) -> Array:
     def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
         if stream is not None:
             raise ValueError("The stream argument to to_device() is not supported")
-        if device == 'cpu':
+        if device == CPU_DEVICE:
             return self
         raise ValueError(f"Unsupported device {device!r}")
 
@@ -1078,11 +1092,11 @@ def dtype(self) -> Dtype:
 
         See its docstring for more information.
         """
-        return self._array.dtype
+        return self._dtype
 
     @property
     def device(self) -> Device:
-        return "cpu"
+        return CPU_DEVICE
 
     # Note: mT is new in array API spec (see matrix_transpose)
     @property

array_api_strict/_creation_functions.py

@@ -12,20 +12,15 @@
         SupportsBufferProtocol,
     )
     from collections.abc import Sequence
-from ._dtypes import _all_dtypes
+from ._dtypes import _DType, _all_dtypes
 
 import numpy as np
 
 
 def _check_valid_dtype(dtype):
     # Note: Only spelling dtypes as the dtype objects is supported.
-
-    # We use this instead of "dtype in _all_dtypes" because the dtype objects
-    # define equality with the sorts of things we want to disallow.
-    for d in (None,) + _all_dtypes:
-        if dtype is d:
-            return
-    raise ValueError("dtype must be one of the supported dtypes")
+    if dtype not in (None,) + _all_dtypes:
+        raise ValueError("dtype must be one of the supported dtypes")
 
 
 def asarray(
@@ -50,10 +45,13 @@ def asarray(
     """
     # _array_object imports in this file are inside the functions to avoid
     # circular imports
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    _np_dtype = None
+    if dtype is not None:
+        _np_dtype = dtype._np_dtype
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
     if copy in (False, np._CopyMode.IF_NEEDED):
         # Note: copy=False is not yet implemented in np.asarray
@@ -62,13 +60,13 @@ def asarray(
         if dtype is not None and obj.dtype != dtype:
             copy = True
         if copy in (True, np._CopyMode.ALWAYS):
-            return Array._new(np.array(obj._array, copy=True, dtype=dtype))
+            return Array._new(np.array(obj._array, copy=True, dtype=_np_dtype))
         return obj
     if dtype is None and isinstance(obj, int) and (obj > 2 ** 64 or obj < -(2 ** 63)):
         # Give a better error message in this case. NumPy would convert this
         # to an object array. TODO: This won't handle large integers in lists.
         raise OverflowError("Integer out of bounds for array dtypes")
-    res = np.asarray(obj, dtype=dtype)
+    res = np.asarray(obj, dtype=_np_dtype)
     return Array._new(res)
 
 
@@ -86,11 +84,13 @@ def arange(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype))
 
 
@@ -105,11 +105,13 @@ def empty(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.empty(shape, dtype=dtype))
 
 
@@ -121,11 +123,13 @@ def empty_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.empty_like(x._array, dtype=dtype))
 
 
@@ -143,11 +147,13 @@ def eye(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype))
 
 
@@ -169,15 +175,17 @@ def full(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
     if isinstance(fill_value, Array) and fill_value.ndim == 0:
         fill_value = fill_value._array
+    if dtype is not None:
+        dtype = dtype._np_dtype
     res = np.full(shape, fill_value, dtype=dtype)
-    if res.dtype not in _all_dtypes:
+    if _DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
         # coerces to one of the acceptable dtypes.
         raise TypeError("Invalid input to full")
@@ -197,13 +205,15 @@ def full_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     res = np.full_like(x._array, fill_value, dtype=dtype)
-    if res.dtype not in _all_dtypes:
+    if _DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
         # coerces to one of the acceptable dtypes.
         raise TypeError("Invalid input to full_like")
@@ -225,11 +235,13 @@ def linspace(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint))
 
 
@@ -264,11 +276,13 @@ def ones(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.ones(shape, dtype=dtype))
 
 
@@ -280,11 +294,13 @@ def ones_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.ones_like(x._array, dtype=dtype))
 
 
@@ -327,11 +343,13 @@ def zeros(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.zeros(shape, dtype=dtype))
 
 
@@ -343,9 +361,11 @@ def zeros_like(
 
     See its docstring for more information.
     """
-    from ._array_object import Array
+    from ._array_object import Array, CPU_DEVICE
 
     _check_valid_dtype(dtype)
-    if device not in ["cpu", None]:
+    if device not in [CPU_DEVICE, None]:
         raise ValueError(f"Unsupported device {device!r}")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     return Array._new(np.zeros_like(x._array, dtype=dtype))

array_api_strict/_data_type_functions.py

@@ -2,6 +2,7 @@
 
 from ._array_object import Array
 from ._dtypes import (
+    _DType,
     _all_dtypes,
     _boolean_dtypes,
     _signed_integer_dtypes,
@@ -27,7 +28,7 @@
 def astype(x: Array, dtype: Dtype, /, *, copy: bool = True) -> Array:
     if not copy and dtype == x.dtype:
         return x
-    return Array._new(x._array.astype(dtype=dtype, copy=copy))
+    return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy))
 
 
 def broadcast_arrays(*arrays: Array) -> List[Array]:
@@ -107,6 +108,8 @@ def finfo(type: Union[Dtype, Array], /) -> finfo_object:
 
     See its docstring for more information.
     """
+    if isinstance(type, _DType):
+        type = type._np_dtype
     fi = np.finfo(type)
     # Note: The types of the float data here are float, whereas in NumPy they
     # are scalars of the corresponding float dtype.
@@ -126,6 +129,8 @@ def iinfo(type: Union[Dtype, Array], /) -> iinfo_object:
 
     See its docstring for more information.
     """
+    if isinstance(type, _DType):
+        type = type._np_dtype
     ii = np.iinfo(type)
     return iinfo_object(ii.bits, ii.max, ii.min, ii.dtype)