More multi-device support

Author: betatim

array_api_strict/_array_object.py

@@ -62,6 +62,7 @@ def __hash__(self):
 
 
 CPU_DEVICE = Device()
+ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"))
 
 _default = object()
 

array_api_strict/_creation_functions.py

@@ -32,11 +32,14 @@ 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 Device
+    from ._array_object import Device, ALL_DEVICES
 
     if device is not None and not isinstance(device, Device):
         raise ValueError(f"Unsupported device {device!r}")
 
+    if device not in ALL_DEVICES:
+        raise ValueError(f"Unsupported device {device!r}")
+
 def asarray(
     obj: Union[
         Array,

array_api_strict/_data_type_functions.py

@@ -37,10 +37,12 @@ def astype(
             _check_device(device)
         else:
             raise TypeError("The device argument to astype requires at least version 2023.12 of the array API")
+    else:
+        device = x.device
 
     if not copy and dtype == x.dtype:
         return x
-    return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy))
+    return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy), device=device)
 
 
 def broadcast_arrays(*arrays: Array) -> List[Array]:

array_api_strict/_fft.py

@@ -14,7 +14,7 @@
     float32,
     complex64,
 )
-from ._array_object import Array, CPU_DEVICE
+from ._array_object import Array, ALL_DEVICES
 from ._data_type_functions import astype
 from ._flags import requires_extension
 
@@ -36,7 +36,7 @@ def fft(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in fft")
-    res = Array._new(np.fft.fft(x._array, n=n, axis=axis, norm=norm))
+    res = Array._new(np.fft.fft(x._array, n=n, axis=axis, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -59,7 +59,7 @@ def ifft(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in ifft")
-    res = Array._new(np.fft.ifft(x._array, n=n, axis=axis, norm=norm))
+    res = Array._new(np.fft.ifft(x._array, n=n, axis=axis, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -82,7 +82,7 @@ def fftn(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in fftn")
-    res = Array._new(np.fft.fftn(x._array, s=s, axes=axes, norm=norm))
+    res = Array._new(np.fft.fftn(x._array, s=s, axes=axes, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -105,7 +105,7 @@ def ifftn(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in ifftn")
-    res = Array._new(np.fft.ifftn(x._array, s=s, axes=axes, norm=norm))
+    res = Array._new(np.fft.ifftn(x._array, s=s, axes=axes, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -128,7 +128,7 @@ def rfft(
     """
     if x.dtype not in _real_floating_dtypes:
         raise TypeError("Only real floating-point dtypes are allowed in rfft")
-    res = Array._new(np.fft.rfft(x._array, n=n, axis=axis, norm=norm))
+    res = Array._new(np.fft.rfft(x._array, n=n, axis=axis, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == float32:
@@ -151,7 +151,7 @@ def irfft(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in irfft")
-    res = Array._new(np.fft.irfft(x._array, n=n, axis=axis, norm=norm))
+    res = Array._new(np.fft.irfft(x._array, n=n, axis=axis, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -174,7 +174,7 @@ def rfftn(
     """
     if x.dtype not in _real_floating_dtypes:
         raise TypeError("Only real floating-point dtypes are allowed in rfftn")
-    res = Array._new(np.fft.rfftn(x._array, s=s, axes=axes, norm=norm))
+    res = Array._new(np.fft.rfftn(x._array, s=s, axes=axes, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == float32:
@@ -197,7 +197,7 @@ def irfftn(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in irfftn")
-    res = Array._new(np.fft.irfftn(x._array, s=s, axes=axes, norm=norm))
+    res = Array._new(np.fft.irfftn(x._array, s=s, axes=axes, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -220,7 +220,7 @@ def hfft(
     """
     if x.dtype not in _complex_floating_dtypes:
         raise TypeError("Only complex floating-point dtypes are allowed in hfft")
-    res = Array._new(np.fft.hfft(x._array, n=n, axis=axis, norm=norm))
+    res = Array._new(np.fft.hfft(x._array, n=n, axis=axis, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == complex64:
@@ -243,7 +243,7 @@ def ihfft(
     """
     if x.dtype not in _real_floating_dtypes:
         raise TypeError("Only real floating-point dtypes are allowed in ihfft")
-    res = Array._new(np.fft.ihfft(x._array, n=n, axis=axis, norm=norm))
+    res = Array._new(np.fft.ihfft(x._array, n=n, axis=axis, norm=norm), device=x.device)
     # Note: np.fft functions improperly upcast float32 and complex64 to
     # complex128
     if x.dtype == float32:
@@ -257,9 +257,9 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Ar
 
     See its docstring for more information.
     """
-    if device not in [CPU_DEVICE, None]:
+    if device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
-    return Array._new(np.fft.fftfreq(n, d=d))
+    return Array._new(np.fft.fftfreq(n, d=d), device=device)
 
 @requires_extension('fft')
 def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
@@ -268,9 +268,9 @@ def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> A
 
     See its docstring for more information.
     """
-    if device not in [CPU_DEVICE, None]:
+    if device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
-    return Array._new(np.fft.rfftfreq(n, d=d))
+    return Array._new(np.fft.rfftfreq(n, d=d), device=device)
 
 @requires_extension('fft')
 def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
@@ -281,7 +281,7 @@ def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
     """
     if x.dtype not in _floating_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in fftshift")
-    return Array._new(np.fft.fftshift(x._array, axes=axes))
+    return Array._new(np.fft.fftshift(x._array, axes=axes), device=x.device)
 
 @requires_extension('fft')
 def ifftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
@@ -292,7 +292,7 @@ def ifftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
     """
     if x.dtype not in _floating_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in ifftshift")
-    return Array._new(np.fft.ifftshift(x._array, axes=axes))
+    return Array._new(np.fft.ifftshift(x._array, axes=axes), device=x.device)
 
 __all__ = [
     "fft",

array_api_strict/tests/test_device_support.py

@@ -0,0 +1,25 @@
+import pytest
+
+import array_api_strict
+
+
+@pytest.mark.parametrize("func_name", ("fft", "ifft", "fftn", "ifftn", "irfft",
+                                  "irfftn", "hfft", "fftshift", "ifftshift"))
+def test_fft_device_support_complex(func_name):
+    func = getattr(array_api_strict.fft, func_name)
+    x = array_api_strict.asarray([1, 2.],
+                                 dtype=array_api_strict.complex64,
+                                 device=array_api_strict.Device("device1"))
+    y = func(x)
+
+    assert x.device == y.device
+
+
+@pytest.mark.parametrize("func_name", ("rfft", "rfftn", "ihfft"))
+def test_fft_device_support_real(func_name):
+    func = getattr(array_api_strict.fft, func_name)
+    x = array_api_strict.asarray([1, 2.],
+                                 device=array_api_strict.Device("device1"))
+    y = func(x)
+
+    assert x.device == y.device