Improvements to device support

[asmeurer]

#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.

[ogrisel]

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

[betatim]

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

[lucyleeow]

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:

1. no mixing of devices
2. 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 )

[ev-br]

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

Two quick notes:

1. 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

2. 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).

[lucascolley]

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

[betatim]

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?

[lucyleeow]

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

[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)