Improve DLPack support for non CPU devices - #233
Open
betatim wants to merge 3 commits into
Open
Conversation
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.
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 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This improves the
from_dlpackbehaviour 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:
Before this PR the assert will fail.
But, now we have a
isinstance(x, Array)branch infrom_dlpack(line 257 ofarray_api_strict/_creation_functions.py). So we could do some special casing in thisifstatement. 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
mainit is possible to "transfer" a device from"device1"to the CPU by usingnp.from_dlpack(x_on_device1). This is explicitly not allowed when usingnp.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.