BUG: Revert dlpack device

[ev-br]

closes gh-219

reverting gh-212 is a bit manual since after landing, things moved around between _array_object.py and _device.py

[ev-br]

Okay, this one is ready from my side. Please do check me, since it's a bit manual and not just a straight revert.

[betatim]

Two small comments.

Ideally we'd have a test that uses pytorch so that we encode that learning about DLPack device IDs into a test (as in: don't mess with them). Can we do that without having to install pytest? Maybe a test that no matter which array-api-strict device is used the device ID is always CPU?

[betatim]

We could keep this test no? It seems like it checks something that is somewhat related to the DLPack device ID but also not related.

I'd keep it (unless I'm missing something)

[ev-br]

An issue is that this test just did not work in 2.5? So if we're just reverting the change made in 2.6, the test is broken; if we also fix from_dlpack, that's not just a revert...

In [1]: import array_api_strict as xp

In [2]: xp.__version__
Out[2]: '2.5'

In [3]: a = xp.arange(3.0, device=xp.Device('device1'))

In [4]: xp.from_dlpack(a).device
Out[4]: array_api_strict.Device('CPU_DEVICE')

---

An attempt at a fix:

$ git diff
diff --git a/array_api_strict/_creation_functions.py b/array_api_strict/_creation_functions.py
index 685044d..5e40771 100644
--- a/array_api_strict/_creation_functions.py
+++ b/array_api_strict/_creation_functions.py
@@ -250,6 +250,8 @@ def from_dlpack(
         _check_device(device)
     else:
         device = None
+        if hasattr(x, "device"):
+            device = x.device
 
     if copy in [_undef, None]:
         # numpy 1.26 does not have the copy= arg

Looks simple and correct, right?

array_api_strict/tests/test_creation_functions.py::test_from_dlpack_default_device FAILED                           [ 96%]
>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> traceback >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

    def test_from_dlpack_default_device():
        x = asarray([1, 2, 3])
        y = from_dlpack(x)
        z = from_dlpack(np.asarray([1, 2, 3]))
>       assert x.device == y.device == z.device == CPU_DEVICE
E       AssertionError: assert array_api_strict.Device('CPU_DEVICE') == 'cpu'
E        +  where array_api_strict.Device('CPU_DEVICE') = Array([1, 2, 3], dtype=array_api_strict.int64).device
E        +  and   'cpu' = Array([1, 2, 3], dtype=array_api_strict.int64, device=cpu).device

So one cannot just reuse the device of x, one has to map the device of x onto the device of array-api-strict! And the device naming is not standardized, so pretty much the only way to map them is via DLPACK enum value, but that's what we reverting here.

TBH I don't see a way out, do you @betatim ?

[betatim]

The reason I thought that this test should "just work" is that both x and y are array-api-strict arrays. But spending two minutes staring at the newly failing test in your comment doesn't really help me understand why your proposed fix shouldn't work.

I think this means that this is more complicated than I think :D So fine to fix this in the future or never or realise that this isn't actually broken

[ev-br]

This one is confusing yes! And I had a head start to stare at this during EuroScipy. The bottom line AFAIU is this. Currently (in 2.5 and 2.6.1):

In [1]: import numpy as np

In [2]: import array_api_strict as xp

In [4]: a = np.ones(3)

In [5]: b = xp.from_dlpack(a)

In [6]: b.device
Out[6]: array_api_strict.Device('CPU_DEVICE')

In [7]: a.device
Out[7]: 'cpu'

The failed fix from #221 (comment) was checking hasattr(a, "device") which was a Yes, and then it attached a.device to b. But a.device is a string 'cpu' not an xp.Device instance.

Note that it currently works mainly by chance: xp.from_dlpack uses the default device unless explicitly provided as a kwarg:

In [8]: x = xp.ones(3, device=xp.Device("device1"))

In [9]: xp.from_dlpack(x).device
Out[9]: array_api_strict.Device('CPU_DEVICE')

So the pickle we're in for xp.from_dlpack(x) is this:

• either we do if hasattr(x, "device"): device=x.device dance --- and then if x is a foreign array, we attach a foreign device to an xp-array;
• or we do what 2.5 was doing, which is device=None means the array-api-strict default device---and then if x is an xp-array on a non-default device, we effectively transfer it to the default "CPU_DEVICE".

Your test was guarding for the latter, but the fix was breaking the former.

---

Now, the 2.5 and 2.6.1 behavior, strictly speaking, is questionable in view of the spec requirement (emphasis mine)

device on which to place the created array. If device is None and x supports DLPack, the output array must be on the same device as x. Default: None.

So which array-api-strict device is the same as NumPy's "cpu" device? And if "device1" is different, then how to tell it is different without somehow mapping the device values between different libraries?

[betatim]

Naively (as an unsuspecting user) the behaviour that the standard describes makes sense. It is what I'd expect to happen.

That the device changes in:

In [8]: x = xp.ones(3, device=xp.Device("device1"))

In [9]: xp.from_dlpack(x).device
Out[9]: array_api_strict.Device('CPU_DEVICE')

seems super weird as a user because we don't change array library in the process.

I think attaching the device from a different array library to the result of from_dlpack is the wrong thing to do. The standard says that how a device is specified is an implementation detail of the library (some use a string, some a class, etc). So attaching a pytorch device to a cupy array won't do anything useful, as none of the cupy code will understand it.

I think the question of "how do you translate a device?" is the key. Can we use the device type and ID of the source array to find the corresponding device in array-api-strict?

In [3]: x = torch.asarray([1.,2,.3])

In [4]: x.__dlpack_device__()
Out[4]: (<DLDeviceType.kDLCPU: 1>, 0)

In [5]: x = torch.asarray([1.,2,.3], device="mps")

In [6]: x.__dlpack_device__()
Out[6]: (<DLDeviceType.kDLMetal: 8>, 0)

So I'd adjust the device handling in

array-api-strict/array_api_strict/_creation_functions.py

Line 234 in f02a5da

def from_dlpack(

to call __dlpack_device__() and see if there is a matching device in array-api-strict if the argument device=None. My intuition would be that __dlpack_device__ returns (1, 0) for Device("CPU_DEVICE") and (1, 1) for Device("device1"). Or something like that. As we learnt we can't change the type because that screws up things, but I think we can use the device ID to tell these otherwise similar devices apart.

[ev-br]

Agreed that the current behaviour is weird! And agreed about the key question, too. Any chance you've cycles to spare to prototype a solution? You seem to have a setup for downstream testing (in scikit-learn or where was that the previous attempt wrought havoc in?).

[ev-br]

In view of #221 (comment), my plan is to wait for a couple of days to see if somebody tells me I miss something simple; and if not, then land this PR and cut a 2.6.1 release.

And then we'll be able to continue crafting dlpack improvements.

[ev-br]

In the interest of unblocking downstream work, let's land this as is, and cut a 2.6.1 bugfix release.
Thanks for the review @betatim , let's continue polishing the dlpack support in a follow-up or several.

[betatim]

Thanks for moving this forward. I was out since Friday, so only just catching up with things again. Looks good to me/I support your decision