ENH: make device2 default to float32 (but still support float64)

This way, "device2" mimics a pytorch CPU device (supports f64,
defaults to f32), and "no_float64" mimics an MPS device (does
not support double precision at all).

While at it, fix the logic in asarray: whether a device does or does not
support a dtype is different from what is the default dtype for this device.
The the decision on the latter should not depend on the former.

Author: ev-br

array_api_strict/_creation_functions.py

@@ -110,23 +110,22 @@ def asarray(
     # numpy default dtype may differ; if so, adjust the dtype
     if dtype is None and device is not None:
         res_dtype = DType(res.dtype)
-        if not device_supports_dtype(device, res_dtype):
-            # The dtype selected by Numpy might not be the default dtype
-            # on this device. If the dtype is not supported by the device we
-            # try to find one that is, aka the default dtype of this device.
-            from ._data_type_functions import isdtype
-            if isdtype(res_dtype, "bool"):
-                target_dtype = DType("bool")
-            elif isdtype(res_dtype, "integral"):
-                target_dtype = get_default_dtypes(device)["integral"]
-            elif isdtype(res_dtype, "real floating"):
-                target_dtype = get_default_dtypes(device)["real floating"]
-            elif isdtype(res_dtype, "complex floating"):
-                target_dtype = get_default_dtypes(device)["complex floating"]
-            else:
-                raise ValueError(f"{res_dtype = } not understood.")
+        # The dtype selected by Numpy might not be the default dtype
+        # on this device. We thus find the default dtype for the dtype "kind", and
+        # cast to the device-appropriate default.
+        from ._data_type_functions import isdtype
+        if isdtype(res_dtype, "bool"):
+            target_dtype = DType("bool")
+        elif isdtype(res_dtype, "integral"):
+            target_dtype = get_default_dtypes(device)["integral"]
+        elif isdtype(res_dtype, "real floating"):
+            target_dtype = get_default_dtypes(device)["real floating"]
+        elif isdtype(res_dtype, "complex floating"):
+            target_dtype = get_default_dtypes(device)["complex floating"]
+        else:
+            raise ValueError(f"{res_dtype = } not understood.")
 
-            res = res.astype(target_dtype._np_dtype)
+        res = res.astype(target_dtype._np_dtype)
 
     return Array._new(res, device=device)
 

array_api_strict/_devices.py

@@ -78,6 +78,15 @@ def check_device(device: Device | None) -> None:
 
 def get_default_dtypes(device: Device | None = None) -> dict[str, DType]:
     if device == NO_FLOAT64_DEVICE:
+        # mimic an MPS device which does not have float64 at all
+        return {
+            "real floating": float32,
+            "complex floating": complex64,
+            "integral": int64,
+            "indexing": int64,
+        }
+    elif device == Device('device2'):
+        # mimic a torch CPU device: support float64 but default to float32
         return {
             "real floating": float32,
             "complex floating": complex64,

array_api_strict/tests/test_creation_functions.py

@@ -22,7 +22,7 @@
     zeros,
     zeros_like,
 )
-from .._dtypes import float32, float64, bool as xp_bool
+from .._dtypes import float32, float64, complex64, bool as xp_bool
 from .._array_object import Array
 from .._devices import CPU_DEVICE, ALL_DEVICES, Device
 from .._info import __array_namespace_info__
@@ -356,6 +356,18 @@ def test_asarray(self):
             asarray(src, device=device)
 
 
+def test_asarray_device_2():
+    # device2 allows float64 but defaults to float32
+    x = asarray([1.0], device=Device('device2'))
+    assert x.dtype == float32
+
+    x = asarray([1j], device=Device('device2'))
+    assert x.dtype == complex64
+
+    y = asarray([1.0], device=Device('device2'), dtype=float64)
+    assert y.dtype == float64
+
+
 @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12'])
 def from_dlpack_2023_12(api_version):
     if api_version != '2022.12':