ENH: asarray handles f32-only device

Author: ev-br

array_api_strict/_creation_functions.py

@@ -106,6 +106,27 @@ def asarray(
         raise OverflowError("Integer out of bounds for array dtypes")
 
     res = np.array(obj, dtype=_np_dtype, copy=copy)
+
+    # 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):
+            # find out the default dtype for the device
+            from ._data_type_functions import isdtype
+            if isdtype(res_dtype, "bool"):
+                targ_dtype = DType("bool")
+            elif isdtype(res_dtype, "integral"):
+                targ_dtype = get_default_dtypes(device)["integral"]
+            elif isdtype(res_dtype, "real floating"):
+                targ_dtype = get_default_dtypes(device)["real floating"]
+            elif isdtype(res_dtype, "complex floating"):
+                targ_dtype = get_default_dtypes(device)["complex floating"]
+            else:
+                raise ValueError(f"{res_dtype = } not understood.")
+            del isdtype
+
+            res = res.astype(targ_dtype._np_dtype)
+
     return Array._new(res, device=device)
 
 

array_api_strict/tests/test_creation_functions.py

@@ -22,7 +22,7 @@
     zeros,
     zeros_like,
 )
-from .._dtypes import float32, float64
+from .._dtypes import float32, float64, bool as xp_bool
 from .._array_object import Array
 from .._devices import CPU_DEVICE, ALL_DEVICES, Device
 from .._info import __array_namespace_info__
@@ -250,6 +250,7 @@ def test_ones_like_etc_correct(self, func):
         device = Device('F32_device')
         b = func(a, device=device)
         assert b.dtype == self.info.default_dtypes(device=device)["real floating"]
+        assert b.device == device
 
     @pytest.mark.parametrize("func", [empty_like, zeros_like, ones_like, _full_like])
     def test_ones_like_etc_incorrect(self, func):
@@ -277,6 +278,7 @@ def test_eye(self):
         device = Device('F32_device')
         a = eye(3, device=device)
         assert a.dtype == self.info.default_dtypes(device=device)["real floating"]
+        assert a.device == device
 
         with pytest.raises((TypeError, ValueError)):
             eye(3, device=device, dtype=float64)
@@ -286,6 +288,7 @@ def test_linspace(self):
 
         a = linspace(1, 10, 11, device=device)
         assert a.dtype == self.info.default_dtypes(device=device)["real floating"]
+        assert a.device == device
 
         a = linspace(1+0j, 10, 11, device=device)
         assert a.dtype == self.info.default_dtypes(device=device)["complex floating"]
@@ -298,18 +301,59 @@ def test_arange(self):
 
         a = arange(0, 10, 1, device=device)
         assert a.dtype == self.info.default_dtypes(device=device)["integral"]
+        assert a.device == device
 
         a = arange(0.0, 10, 1, device=device)
         assert a.dtype == self.info.default_dtypes(device=device)["real floating"]
+        assert a.device == device
 
         with pytest.raises((TypeError, ValueError)):
             arange(0, 10, 1, device=device, dtype=float64)
 
         with pytest.raises((TypeError, ValueError)):
             arange(0.0, 10, 1, device=device, dtype=float64)
 
-# TODO:
-# def asarray(
+    def test_asarray(self):
+        device = Device('F32_device')
+
+        ### asarray(python_object)
+        for x in (True, [False,]):
+            arr = asarray(x, device=device)
+            assert arr.dtype == xp_bool
+            assert arr.device == device
+
+        for x in [1, [1,]]:
+            arr = asarray(x, device=device)
+            assert arr.dtype == self.info.default_dtypes(device=device)['integral']
+            assert arr.device == device
+
+        for x in [1.0, [1.0,]]:
+            arr = asarray(x, device=device)
+            assert arr.dtype == self.info.default_dtypes(device=device)['real floating']
+            assert arr.device == device
+
+        for x in [1j, [1j,]]:
+            arr = asarray(x, device=device)
+            assert arr.dtype == self.info.default_dtypes(device=device)['complex floating']
+            assert arr.device == device
+
+        # asarray(python_object, dtype=unsupported_by_device)
+        with pytest.raises(ValueError, match="Device"):
+            asarray(1, dtype=float64, device=device)
+
+        ### asarray(array)
+
+        # compatible dtypes, device transfer
+        src = asarray(1, dtype=float32, device=Device('device1'))
+        dst = asarray(src, device=device)
+        assert dst.device == device
+        assert dst.dtype == float32
+
+        # incompatible dtypes, device transfer
+        src = asarray(1, dtype=float64, device=Device('device1'))
+
+        with pytest.raises(ValueError, match="Device"):
+            asarray(src, device=device)
 
 
 @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12'])