squash

Author: ev-br

array_api_strict/_array_object.py

@@ -49,7 +49,7 @@ class Device:
     __slots__ = ("_device", "__weakref__")
 
     def __init__(self, device: str = "CPU_DEVICE"):
-        if device not in ("CPU_DEVICE", "device1", "device2"):
+        if device not in ("CPU_DEVICE", "device1", "device2", "F32_device"):
             raise ValueError(f"The device '{device}' is not a valid choice.")
         self._device = device
 
@@ -66,7 +66,7 @@ def __hash__(self) -> int:
 
 
 CPU_DEVICE = Device()
-ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"))
+ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), Device("F32_device"))
 
 
 class Array:

array_api_strict/tests/test_device_support.py

@@ -36,3 +36,20 @@ def test_fft_device_support_real(func_name):
     y = func(x)
 
     assert x.device == y.device
+
+
+class TestF32Device:
+    @pytest.mark.parametrize("dtype_str", ["float64", "complex128"])
+    def test_f64_raises(self, dtype_str):
+        f32_device = array_api_strict.Device("F32_device")
+        dtype = getattr(array_api_strict, dtype_str)
+        with pytest.raises(ValueError):
+            array_api_strict.arange(3, device=f32_device, dtype=dtype)
+
+    def test_info_no_f64(self):
+        f32_device = array_api_strict.Device("F32_device")
+
+        info = array_api_strict.__array_namespace_info__()
+        all_dtypes = info.dtypes(device=f32_device)
+        assert "float64" not in all_dtypes
+        assert "complex128" not in all_dtypes