ENH: add a "no_x64" device which only supports 32-bit ints and floats

This mimics the JAX's default (and the name follows the
JAX `jax_enable_x64` flag name).

Author: ev-br

array_api_strict/_devices.py

@@ -1,13 +1,13 @@
 from typing import Final
 
 from ._dtypes import (
-    DType, float32, float64, complex64, complex128, int64,
+    DType, float32, float64, complex64, complex128, int64, uint64, int32,
     _all_dtypes, _boolean_dtypes, _signed_integer_dtypes,
     _unsigned_integer_dtypes, _integer_dtypes, _real_floating_dtypes,
     _complex_floating_dtypes, _numeric_dtypes
 )
 
-_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64")
+_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64", "no_x64")
 
 class Device:
     _device: Final[str]
@@ -32,8 +32,11 @@ def __hash__(self) -> int:
 
 CPU_DEVICE = Device()
 NO_FLOAT64_DEVICE = Device("no_float64")
+NO_X64_DEVICE = Device("no_x64")
 
-ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE)
+ALL_DEVICES = (
+    CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE, NO_X64_DEVICE
+)
 
 
 def check_device(device: Device | None) -> None:
@@ -55,6 +58,14 @@ def get_default_dtypes(device: Device | None = None) -> dict[str, DType]:
             "integral": int64,
             "indexing": int64,
         }
+    elif device == NO_X64_DEVICE:
+        # mimic JAX default: no float64, no int64
+        return {
+            "real floating": float32,
+            "complex floating": complex64,
+            "integral": int32,
+            "indexing": int32,
+        }
     elif device == Device('device2'):
         # mimic a torch CPU device: support float64 but default to float32
         return {
@@ -77,6 +88,9 @@ def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool:
     # Device("no_float64") supports all dtypes except float64 and complex128
     if device == NO_FLOAT64_DEVICE:
         return dtype not in (float64, complex128)
+    if device == NO_X64_DEVICE:
+        # "no_x64" only supports 32-bit ints and floats
+        return dtype not in (float64, complex128, int64, uint64)
 
     # All other devices support all dtypes
     return True

array_api_strict/tests/test_creation_functions.py

@@ -22,7 +22,7 @@
     zeros,
     zeros_like,
 )
-from .._dtypes import float32, float64, complex64, bool as xp_bool
+from .._dtypes import float32, float64, complex64, int32, int64, bool as xp_bool
 from .._array_object import Array
 from .._devices import CPU_DEVICE, ALL_DEVICES, Device
 from .._info import __array_namespace_info__
@@ -368,6 +368,17 @@ def test_asarray_device_2():
     assert y.dtype == float64
 
 
+def test_asarray_no_x64_device():
+    x = asarray(3, device=Device("no_x64"))
+    assert x.dtype == int32
+
+    with pytest.raises(ValueError):
+        asarray(3, device=Device("no_x64"), dtype=int64)
+
+    y = zeros_like(ones(3, dtype=int32), device=Device("no_x64"))
+    assert y.dtype == int32
+
+
 @pytest.mark.parametrize("api_version", ['2021.12', '2022.12', '2023.12'])
 def from_dlpack_2023_12(api_version):
     if api_version != '2022.12':