ENH: info.default_dtypes: make device-aware

Author: ev-br

array_api_strict/_devices.py

@@ -1,7 +1,7 @@
 from typing import Final
 
-from ._dtypes import DType, float64, complex128
 from ._dtypes import (
+    DType, float32, float64, complex64, complex128, int64,
     _all_dtypes, _boolean_dtypes, _signed_integer_dtypes,
     _unsigned_integer_dtypes, _integer_dtypes, _real_floating_dtypes,
     _complex_floating_dtypes, _numeric_dtypes
@@ -50,6 +50,22 @@ def check_device(device: Device | None) -> None:
 
 # Helpers for device-specific dtype support
 
+def get_default_dtypes(device: Device | None = None) -> dict[str, Device]:
+    if device == _F32_DEVICE:
+        return {
+            "real floating": float32,
+            "complex floating": complex64,
+            "integral": int64,
+            "indexing": int64,
+        }
+    else:
+        return {
+            "real floating": float64,
+            "complex floating": complex128,
+            "integral": int64,
+            "indexing": int64,
+        }
+
 
 def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool:
     """True if `device` supports `dtype`, False otherwise."""

array_api_strict/_info.py

@@ -1,8 +1,7 @@
 import numpy as np
 
-from . import _dtypes as dt
 from . import _devices
-from ._array_object import ALL_DEVICES, CPU_DEVICE, Device
+from ._devices import ALL_DEVICES, CPU_DEVICE, Device
 from ._flags import get_array_api_strict_flags, requires_api_version
 from ._typing import Capabilities, DataTypes, DefaultDataTypes
 
@@ -41,12 +40,7 @@ def default_dtypes(
         *,
         device: Device | None = None,
     ) -> DefaultDataTypes:
-        return {
-            "real floating": dt.float64,
-            "complex floating": dt.complex128,
-            "integral": dt.int64,
-            "indexing": dt.int64,
-        }
+        return _devices.get_default_dtypes(device)
 
     @requires_api_version('2023.12')
     def dtypes(

array_api_strict/tests/test_device_support.py

@@ -53,3 +53,16 @@ def test_info_no_f64(self):
         all_dtypes = info.dtypes(device=f32_device)
         assert "float64" not in all_dtypes
         assert "complex128" not in all_dtypes
+
+    def test_info_default_dtypes(self):
+        f32_device = xp.Device("F32_device")
+        info = xp.__array_namespace_info__()
+        defaults = info.default_dtypes(device=f32_device)
+        assert defaults["real floating"] == xp.float32
+        assert defaults["complex floating"] == xp.complex64
+
+        cpu_device = xp.Device()
+        info = xp.__array_namespace_info__()
+        defaults = info.default_dtypes(device=cpu_device)
+        assert defaults["real floating"] == xp.float64
+        assert defaults["complex floating"] == xp.complex128