Use separate wrapped dtype objects in numpy.array_api

This way there is no ambiguity about the fact the non-portability of NumPy
dtype behavior, or the fact that NumPy dtypes are not necessarily allowed as
dtypes for non-NumPy array APIs.

Fixes #23883

Original NumPy Commit: 13ab654e46110221b6388aaad606a3625f43db5a

Author: asmeurer

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,
@@ -81,11 +82,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
@@ -107,7 +110,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)
@@ -182,6 +185,7 @@ 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):
@@ -195,7 +199,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"
@@ -221,7 +225,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]:
@@ -331,7 +335,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)
@@ -1084,7 +1090,7 @@ def dtype(self) -> Dtype:
 
         See its docstring for more information.
         """
-        return self._array.dtype
+        return self._dtype
 
     @property
     def device(self) -> Device:

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 not dtype in (None,) + _all_dtypes:
+        raise ValueError("dtype must be one of the supported dtypes")
 
 
 def asarray(
@@ -68,6 +63,8 @@ def asarray(
         # 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")
+    if dtype is not None:
+        dtype = dtype._np_dtype
     res = np.asarray(obj, dtype=dtype)
     return Array._new(res)
 
@@ -91,6 +88,8 @@ def arange(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -110,6 +109,8 @@ def empty(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -126,6 +127,8 @@ def empty_like(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -148,6 +151,8 @@ def eye(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -176,8 +181,10 @@ def full(
         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")
@@ -202,8 +209,10 @@ def full_like(
     _check_valid_dtype(dtype)
     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")
@@ -230,6 +239,8 @@ def linspace(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -269,6 +280,8 @@ def ones(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -285,6 +298,8 @@ def ones_like(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -332,6 +347,8 @@ def zeros(
     _check_valid_dtype(dtype)
     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))
 
 
@@ -348,4 +365,6 @@ def zeros_like(
     _check_valid_dtype(dtype)
     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)
 

array_api_strict/_dtypes.py

@@ -1,21 +1,42 @@
 import numpy as np
 
-# Note: we use dtype objects instead of dtype classes. The spec does not
-# require any behavior on dtypes other than equality.
-int8 = np.dtype("int8")
-int16 = np.dtype("int16")
-int32 = np.dtype("int32")
-int64 = np.dtype("int64")
-uint8 = np.dtype("uint8")
-uint16 = np.dtype("uint16")
-uint32 = np.dtype("uint32")
-uint64 = np.dtype("uint64")
-float32 = np.dtype("float32")
-float64 = np.dtype("float64")
-complex64 = np.dtype("complex64")
-complex128 = np.dtype("complex128")
+# Note: we wrap the NumPy dtype objects in a bare class, so that none of the
+# additional methods and behaviors of NumPy dtype objects are exposed.
+
+class _DType:
+    def __init__(self, np_dtype):
+        np_dtype = np.dtype(np_dtype)
+        self._np_dtype = np_dtype
+
+    def __repr__(self):
+        return f"np.array_api.{self._np_dtype.name}"
+
+    def __eq__(self, other):
+        if not isinstance(other, _DType):
+            return NotImplemented
+        return self._np_dtype == other._np_dtype
+
+    def __hash__(self):
+        # Note: this is not strictly required
+        # (https://github.com/data-apis/array-api/issues/582), but makes the
+        # dtype objects much easier to work with here and elsewhere if they
+        # can be used as dict keys.
+        return hash(self._np_dtype)
+
+int8 = _DType("int8")
+int16 = _DType("int16")
+int32 = _DType("int32")
+int64 = _DType("int64")
+uint8 = _DType("uint8")
+uint16 = _DType("uint16")
+uint32 = _DType("uint32")
+uint64 = _DType("uint64")
+float32 = _DType("float32")
+float64 = _DType("float64")
+complex64 = _DType("complex64")
+complex128 = _DType("complex128")
 # Note: This name is changed
-bool = np.dtype("bool")
+bool = _DType("bool")
 
 _all_dtypes = (
     int8,

array_api_strict/_manipulation_functions.py

@@ -20,7 +20,7 @@ def concat(
     # (no for scalars with axis=None, no cross-kind casting)
     dtype = result_type(*arrays)
     arrays = tuple(a._array for a in arrays)
-    return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype))
+    return Array._new(np.concatenate(arrays, axis=axis, dtype=dtype._np_dtype))
 
 
 def expand_dims(x: Array, /, *, axis: int) -> Array:
@@ -53,8 +53,8 @@ def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
 
 
 # Note: the optional argument is called 'shape', not 'newshape'
-def reshape(x: Array, 
-            /, 
+def reshape(x: Array,
+            /,
             shape: Tuple[int, ...],
             *,
             copy: Optional[Bool] = None) -> Array:

array_api_strict/_statistical_functions.py

@@ -67,9 +67,11 @@ def prod(
     # special-case it here
     if dtype is None:
         if x.dtype == float32:
-            dtype = float64
+            dtype = np.float64
         elif x.dtype == complex64:
-            dtype = complex128
+            dtype = np.complex128
+    else:
+        dtype = dtype._np_dtype
     return Array._new(np.prod(x._array, dtype=dtype, axis=axis, keepdims=keepdims))
 
 
@@ -102,9 +104,11 @@ def sum(
     # special-case it here
     if dtype is None:
         if x.dtype == float32:
-            dtype = float64
+            dtype = np.float64
         elif x.dtype == complex64:
-            dtype = complex128
+            dtype = np.complex128
+    else:
+        dtype = dtype._np_dtype
     return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims))
 
 

array_api_strict/linalg.py

@@ -8,6 +8,7 @@
     complex64,
     complex128
 )
+from ._data_type_functions import finfo
 from ._manipulation_functions import reshape
 from ._array_object import Array
 
@@ -204,7 +205,7 @@ def matrix_rank(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> A
         raise np.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional")
     S = np.linalg.svd(x._array, compute_uv=False)
     if rtol is None:
-        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * np.finfo(S.dtype).eps
+        tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * finfo(S.dtype).eps
     else:
         if isinstance(rtol, Array):
             rtol = rtol._array
@@ -254,7 +255,7 @@ def pinv(x: Array, /, *, rtol: Optional[Union[float, Array]] = None) -> Array:
     # Note: this is different from np.linalg.pinv, which does not multiply the
     # default tolerance by max(M, N).
     if rtol is None:
-        rtol = max(x.shape[-2:]) * np.finfo(x.dtype).eps
+        rtol = max(x.shape[-2:]) * finfo(x.dtype).eps
     return Array._new(np.linalg.pinv(x._array, rcond=rtol))
 
 def qr(x: Array, /, *, mode: Literal['reduced', 'complete'] = 'reduced') -> QRResult:
@@ -384,9 +385,11 @@ def trace(x: Array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> Arr
     # _statistical_functions.py)
     if dtype is None:
         if x.dtype == float32:
-            dtype = float64
+            dtype = np.float64
         elif x.dtype == complex64:
-            dtype = complex128
+            dtype = np.complex128
+    else:
+        dtype = dtype._np_dtype
     # Note: trace always operates on the last two axes, whereas np.trace
     # operates on the first two axes by default
     return Array._new(np.asarray(np.trace(x._array, offset=offset, axis1=-2, axis2=-1, dtype=dtype)))