From 0be49851c4e09fbdd09046493cf9087e9104f87f Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sun, 12 Jul 2026 11:55:01 +0200
Subject: [PATCH] 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).
---
 array_api_strict/_devices.py                  | 20 ++++++++++++++++---
 .../tests/test_creation_functions.py          | 13 +++++++++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/array_api_strict/_devices.py b/array_api_strict/_devices.py
index c48db2a..0d47359 100644
--- a/array_api_strict/_devices.py
+++ b/array_api_strict/_devices.py
@@ -2,13 +2,13 @@
 from enum import IntEnum
 
 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]
@@ -33,8 +33,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
+)
 
 class DLDeviceType(IntEnum):
     kDLCPU = 1
@@ -85,6 +88,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 {
@@ -107,6 +118,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
diff --git a/array_api_strict/tests/test_creation_functions.py b/array_api_strict/tests/test_creation_functions.py
index 06c5d87..366ca25 100644
--- a/array_api_strict/tests/test_creation_functions.py
+++ b/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__
@@ -375,6 +375,17 @@ def test_asarray_device_1():
     assert y.dtype == float32
 
 
+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':