Skip to content

Improve DLPack support for non CPU devices - #233

Open
betatim wants to merge 3 commits into
data-apis:mainfrom
betatim:match-from_dlpack_device
Open

Improve DLPack support for non CPU devices#233
betatim wants to merge 3 commits into
data-apis:mainfrom
betatim:match-from_dlpack_device

Conversation

@betatim

@betatim betatim commented Jul 29, 2026

Copy link
Copy Markdown
Member

This improves the from_dlpack behaviour for round-tripping array-api-strict arrays and raises an error for arrays that are not on the CPU device.

This is a follow up to #219 and #221

We use the device ID to tell the difference between the different "devices" present in array-api-strict. This seems to work/not bother libraries like torch or numpy. The device type is always CPU. In #212 we tried using different device type IDs, but that turns out not to work. I think the standard says (or suggests?) that for the CPU device the only device ID that makes sense is zero. In practice it seems every library ignores this field for CPU devices. Originally I thought we needed to use different device IDs to be able to make the following code work:

x = asarray([1, 2, 3], device=device)
y = from_dlpack(x)
assert y.device == device

Before this PR the assert will fail.

But, now we have a isinstance(x, Array) branch in from_dlpack (line 257 of array_api_strict/_creation_functions.py). So we could do some special casing in this if statement. That way we don't have to report IDs that are not part of the standard. The downside is that __dlpack_device__ would report the same for all the devices in array-api-strict. I can't make up my mind which option I prefer.

On main it is possible to "transfer" a device from "device1" to the CPU by using np.from_dlpack(x_on_device1). This is explicitly not allowed when using np.asarray(). This PR adjusts the DLPack behaviour to match the __buffer__ based behaviour.

This also installs torch for one of the CI jobs to test DLpack'ing arrays to and from a different library.

betatim added 2 commits July 29, 2026 13:58
This improves the `from_dlpack` behaviour for round-tripping
array-api-strict arrays and raises an error for arrays that are not on
the CPU device.
The CPU device can not really have multiple devices, and we don't need
to be able to tell devices apart. So this removes the "not spec
compliant" code.
@betatim

betatim commented Jul 29, 2026

Copy link
Copy Markdown
Member Author

To illustrate my indecision, the second commit removes the special "use different device IDs for different logical devices". I think it is cleaner this way

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant