Add "multi device" support

Having more than one device is useful during testing to allow you to
find bugs related to how arrays on different devices are handled.

Author: betatim

array_api_strict/__init__.py

@@ -309,6 +309,9 @@
 
 __all__ += ["all", "any"]
 
+from ._array_object import Device
+__all__ += ["Device"]
+
 # Helper functions that are not part of the standard
 
 from ._flags import (

array_api_strict/_array_object.py

@@ -43,13 +43,17 @@
 
 import numpy as np
 
-# Placeholder object to represent the "cpu" device (the only device NumPy
-# supports).
-class _cpu_device:
+class Device:
+    def __init__(self, device="CPU_DEVICE"):
+        self._device = device
+
     def __repr__(self):
-        return "CPU_DEVICE"
+        return f"Device('{self._device}')"
+
+    def __eq__(self, other):
+        return self._device == other._device
 
-CPU_DEVICE = _cpu_device()
+CPU_DEVICE = Device()
 
 _default = object()
 
@@ -73,7 +77,7 @@ class Array:
     # Use a custom constructor instead of __init__, as manually initializing
     # this class is not supported API.
     @classmethod
-    def _new(cls, x, /):
+    def _new(cls, x, /, device=CPU_DEVICE):
         """
         This is a private method for initializing the array API Array
         object.
@@ -95,6 +99,9 @@ def _new(cls, x, /):
             )
         obj._array = x
         obj._dtype = _dtype
+        if device is None:
+            device = CPU_DEVICE
+        obj._device = device
         return obj
 
     # Prevent Array() from working
@@ -134,6 +141,8 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
         will be present in other implementations.
 
         """
+        if self._device != CPU_DEVICE:
+            raise RuntimeError(f"Can not convert array on the '{self._device}' device to a Numpy array.")
         # copy keyword is new in 2.0.0; for older versions don't use it
         # retry without that keyword.
         if np.__version__[0] < '2':
@@ -1154,8 +1163,11 @@ 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_DEVICE:
+        if device == self._device:
             return self
+        elif isinstance(device, Device):
+            arr = np.asarray(self._array, copy=True)
+            return self.__class__._new(arr, device=device)
         raise ValueError(f"Unsupported device {device!r}")
 
     @property
@@ -1169,7 +1181,7 @@ def dtype(self) -> Dtype:
 
     @property
     def device(self) -> Device:
-        return CPU_DEVICE
+        return self._device
 
     # Note: mT is new in array API spec (see matrix_transpose)
     @property

array_api_strict/_creation_functions.py

@@ -32,9 +32,9 @@ def _supports_buffer_protocol(obj):
 def _check_device(device):
     # _array_object imports in this file are inside the functions to avoid
     # circular imports
-    from ._array_object import CPU_DEVICE
+    from ._array_object import Device
 
-    if device not in [CPU_DEVICE, None]:
+    if device is not None and not isinstance(device, Device):
         raise ValueError(f"Unsupported device {device!r}")
 
 def asarray(
@@ -79,7 +79,7 @@ def asarray(
                 return Array._new(new_array)
             elif _supports_buffer_protocol(obj):
                 # Buffer protocol will always support no-copy
-                return Array._new(np.array(obj, copy=copy, dtype=_np_dtype))
+                return Array._new(np.array(obj, copy=copy, dtype=_np_dtype), device=device)
             else:
                 # No-copy is unsupported for Python built-in types.
                 raise ValueError("Unable to avoid copy while creating an array from given object.")
@@ -89,13 +89,13 @@ def asarray(
             copy = False
 
     if isinstance(obj, Array):
-        return Array._new(np.array(obj._array, copy=copy, dtype=_np_dtype))
+        return Array._new(np.array(obj._array, copy=copy, dtype=_np_dtype), device=device)
     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.array(obj, dtype=_np_dtype, copy=copy)
-    return Array._new(res)
+    return Array._new(res, device=device)
 
 
 def arange(
@@ -119,7 +119,7 @@ def arange(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype))
+    return Array._new(np.arange(start, stop=stop, step=step, dtype=dtype), device=device)
 
 
 def empty(
@@ -140,7 +140,7 @@ def empty(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.empty(shape, dtype=dtype))
+    return Array._new(np.empty(shape, dtype=dtype), device=device)
 
 
 def empty_like(
@@ -158,7 +158,7 @@ def empty_like(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.empty_like(x._array, dtype=dtype))
+    return Array._new(np.empty_like(x._array, dtype=dtype), device=device)
 
 
 def eye(
@@ -182,7 +182,7 @@ def eye(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype))
+    return Array._new(np.eye(n_rows, M=n_cols, k=k, dtype=dtype), device=device)
 
 
 _default = object()
@@ -237,7 +237,7 @@ def full(
         # 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")
-    return Array._new(res)
+    return Array._new(res, device=device)
 
 
 def full_like(
@@ -265,7 +265,7 @@ def full_like(
         # 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")
-    return Array._new(res)
+    return Array._new(res, device=device)
 
 
 def linspace(
@@ -290,7 +290,7 @@ def linspace(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint))
+    return Array._new(np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint), device=device)
 
 
 def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
@@ -308,7 +308,7 @@ def meshgrid(*arrays: Array, indexing: str = "xy") -> List[Array]:
         raise ValueError("meshgrid inputs must all have the same dtype")
 
     return [
-        Array._new(array)
+        Array._new(array, device=device)
         for array in np.meshgrid(*[a._array for a in arrays], indexing=indexing)
     ]
 
@@ -331,7 +331,7 @@ def ones(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.ones(shape, dtype=dtype))
+    return Array._new(np.ones(shape, dtype=dtype), device=device)
 
 
 def ones_like(
@@ -349,7 +349,7 @@ def ones_like(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.ones_like(x._array, dtype=dtype))
+    return Array._new(np.ones_like(x._array, dtype=dtype), device=device)
 
 
 def tril(x: Array, /, *, k: int = 0) -> Array:
@@ -377,7 +377,7 @@ def triu(x: Array, /, *, k: int = 0) -> Array:
     if x.ndim < 2:
         # Note: Unlike np.triu, x must be at least 2-D
         raise ValueError("x must be at least 2-dimensional for triu")
-    return Array._new(np.triu(x._array, k=k))
+    return Array._new(np.triu(x._array, k=k), device=device)
 
 
 def zeros(
@@ -398,7 +398,7 @@ def zeros(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.zeros(shape, dtype=dtype))
+    return Array._new(np.zeros(shape, dtype=dtype), device=device)
 
 
 def zeros_like(
@@ -416,4 +416,4 @@ def zeros_like(
 
     if dtype is not None:
         dtype = dtype._np_dtype
-    return Array._new(np.zeros_like(x._array, dtype=dtype))
+    return Array._new(np.zeros_like(x._array, dtype=dtype), device=device)

array_api_strict/_typing.py

@@ -27,7 +27,7 @@
     Protocol,
 )
 
-from ._array_object import Array, _cpu_device
+from ._array_object import Array, _device
 from ._dtypes import _DType
 
 _T_co = TypeVar("_T_co", covariant=True)
@@ -37,7 +37,7 @@ def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 
-Device = _cpu_device
+Device = _device
 
 Dtype = _DType
 

array_api_strict/tests/test_array_object.py

@@ -319,7 +319,7 @@ def test_python_scalar_construtors():
 def test_device_property():
     a = ones((3, 4))
     assert a.device == CPU_DEVICE
-    assert a.device != 'cpu'
+    assert not isinstance(a.device, str)
 
     assert all(equal(a.to_device(CPU_DEVICE), a))
     assert_raises(ValueError, lambda: a.to_device('cpu'))