Improvements to device support

#59 added support for basic devices. Some improvements that could be made:

Factor out the device checking logic (into a decorator?)
Add devices that don't support certain dtypes (non-support for float64/complex128 being the most common)
Allow the user to manually create custom "devices" with certain properties (probably via the flags API).
Better testing for device support. This ideally should go in array-api-tests, but currently it doesn't test devices at all, so it would be useful to have some basic tests here.

I tried to run the scikit-learn tests with array-api-strict on the non-default device and I got some failures that do not happen with PyTorch and non-default device. So I suspect there is something fishy happening with the new devices in array-api-strict but I have not investigated the root cause yet.

scikit-learn/scikit-learn#30090

I'll take a look/investigate that failure further.

@ogrisel

Commenting here because it partially relates OP but directly related to @ogrisel 's comment #70 (comment)

Looking at: https://data-apis.org/array-api/latest/design_topics/device_support.html#device-support , the spec is very light in this area. It only specifies support for local control over data allocation target device, and syntax for device assignment. I agree that it makes sense that strict does not allow mixing of devices, but there are 2 different options here:

no mixing of devices
only allow the default device (note NumPy itself only allows a default device)

Current behaviour seems a bit conflicting (?), this works:

from array_api_strict import asarray, Device

asarray([1,2,3], device=Device('device1'))

this fails:

asarray([asarray([1,2], device=Device('device1')), asarray([1,2], device=Device('device1'))])

due to this check (in __array__):

array-api-strict/array_api_strict/_array_object.py

Lines 166 to 167 in a8f567a

    
           if self._device != CPU_DEVICE: 
        
               raise RuntimeError(f"Can not convert array on the '{self._device}' device to a Numpy array.")

this seems to be the only place where such a check is happening. And it's probably the cause of the behaviour noticed by @ogrisel where some tests on non-default devices are failing on strict but not with Torch and non-default device.

(I think we should agree on a behaviour and make it consistent)

@ev-br

(ref: noticed while working on #134 )

@lucyleeow

First and foremost, working on device support is very welcome @lucyleeow !

Two quick notes:

Neither dunder array nor nested asarrays are standard-compliant. asarray is not strict enough #118 tracks the latter in array-api-compat and I just sent a PR to potentially clarify the spec, docs: clarify that behavior when provided a nested sequence to asarray is unspecified array-api#917
My gut feeling is that allowing multiple dummy devices in array-api-strict is more useful than only allowing a single CPU_DEVICE. If only to simplify testing device support in array-api-tests (the last bullet in the OP).

Agreed, having different devices is useful for testing.

asarray([asarray([1,2], device=Device('device1')), asarray([1,2], device=Device('device1'))])

is more idiomatically written as

xp.stack((xp.asarray([1, 2], device=Device('device1')), xp.asarray([1, 2], device=Device('device1')))

which should not fail

The idea behind the "multi device support" in array-api-strict is to help projects test the behaviour of using more than one device, without needing actual special hardware (like a CUDA GPU).

The decision I made was to allow arrays on the default device (CPU_DEVICE) to be converted back to Numpy arrays. The reasoning being that it was already possible (don't introduce breaking changes) and that there are libraries that allow this (so it would be useful to test this). The "non default devices" don't allow conversion to Numpy arrays because there are libraries that don't allow you to do that and there was no existing behaviour so we could do what we wanted.

I think the error you get when you write asarray([asarray([1,2], device=Device('device1')), asarray([1,2], device=Device('device1'))]) is a bit weird/hard to figure out. At least I was puzzled by it. Maybe we can improve that to help users and point them to what Lucas suggested as the way to write this? Though I was also not sure what the intended outcome of asarray([asarray([1,2], device=Device('device1')), asarray([1,2], device=Device('device1'))]) is, an array on the default device or on device1?

@betatim

Thanks for the background @betatim, I think your decisions all sound reasonable.

I'm afraid the context is very boring, while working on #134, I wanted to amend test_indexing_arrays to also test device support. That test unfortunately has:

array-api-strict/array_api_strict/tests/test_array_object.py

Line 112 in a8f567a

a_idx_loop = asarray([a[idx[i]] for i in range(idx.shape[0])])

which I think we now agree should have been written xp.stack([a[idx[i]] for i in range(idx.shape[0])])

I agree that the error is confusing. Indeed, this was also the error returned by bug #134 , though that one is an actual bug to be fixed.

I tried to run the scikit-learn tests with array-api-strict on the non-default device and I got some failures that do not happen with PyTorch and non-default device.

I am now interested in what test failed with array-api-strict but passed with PyTorch? @ogrisel

Another thing that could be useful: make it possible to configure an array_api_strict device object that can perform computation on xp.float32 data but would raise an exception when attempting to call functions that need to allocate a device array with xp.float64 dtype.

This is the case for torch with the "mps" device on macOS. However, it's not easy to replicate for people who do not have access to recent Apple hardware. For reference here are the errors one gets with torch via array API compat on MPS:

>>> import array_api_compat.torch as xp
>>> import numpy as np
>>> xp.asarray(np.ones(3, device="cpu", dtype=xp.float64), device="mps")
Traceback (most recent call last):
  Cell In[4], line 1
    xp.asarray(np.ones(3), device="mps")
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/array_api_compat/torch/_aliases.py:205 in asarray
    return torch.asarray(obj, dtype=dtype, device=device, copy=copy, **kwargs)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

>>> xp.from_dlpack(xp.ones(3, device="cpu", dtype=xp.float64), device="mps")
Traceback (most recent call last):
  Cell In[5], line 1
    xp.from_dlpack(xp.ones(3, device="cpu", dtype=xp.float64), device="mps")
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/torch/utils/dlpack.py:220 in from_dlpack
    tensor = tensor.to(device)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

>>> xp.astype(xp.ones(3, device="mps"), xp.float64)
Traceback (most recent call last):
  Cell In[6], line 1
    xp.astype(xp.ones(3, device="mps"), xp.float64)
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/array_api_compat/torch/_aliases.py:726 in astype
    return x.to(dtype=dtype, copy=copy)
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

Ideally, this limitation should be documented in __array_namespace__info__.

EDIT: I just realized that this need was already expressed in the description of the issue:

Add devices that don't support certain dtypes (non-support for float64/complex128 being the most common)

So the requests seem to converge on essentially #64. Here's a WIP/POC enhancement, #206

ISTM it'd be most helpful to split the discussion into two streams:

array-api-strict specific details in ENH: add a float32-only device #206 (the current state is fairly rough, comments / test-driving would be very helpful indeed)
[RFC]: Updates for float16/bfloat16 and for dtypes that are lacking full support in libraries array-api#998 starts an RFC for updating the spec

Another related issue:

Raise on implicit device transfer in __setitem__ #207

array-api-strict==2.6 (released today, 09 Jul 2026) ships with gh-207 (no implicit device tranfers in __setitem__) and gh-206 (a new Device("no_float64") device).

array-api-strict==2.6.1 adds a Device("no_x64") to mimic the default JAX behaviour: only support single precision floats and integers of <=32 bit width.

Navigation Menu

Repository navigation

Improvements to device support #70

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	if self._device != CPU_DEVICE:
	raise RuntimeError(f"Can not convert array on the '{self._device}' device to a Numpy array.")

Improvements to device support #70

Description

Activity

ogrisel commented on Oct 21, 2024

betatim commented on Oct 21, 2024

lucyleeow commented on Mar 28, 2025

ev-br commented on Mar 28, 2025

lucascolley commented on Mar 28, 2025

betatim commented on Mar 31, 2025

lucyleeow commented on Mar 31, 2025

ogrisel commented on Apr 22, 2026

ev-br commented on Apr 24, 2026

ogrisel commented on Apr 30, 2026

ev-br commented on Jul 9, 2026

ev-br commented on Jul 27, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions