ENH: config manager to record all device transfers

[lucascolley]

EDIT: see #168 (comment)

Related to gh-70 but separate from the main discussion there.

This issue follows up on a discussion I had with @ogrisel over dinner the other day, where Olivier suggested that it might be useful for array-api-strict to have a mode which disallows changing the device of an array.

Right now, by default, mixing devices in a single function call is disallowed. However, changing the device of an array, in to_device and array creation functions, is allowed, as far as I can tell from skimming the code.

This may be a sensible default (e.g. otherwise to_device is a bit useless), however it may make sense to be able to override this behaviour, if a consuming library wants to ensure a single device is maintained throughout a test. I believe this is of interest to scikit-learn.

cc @betatim from previous device-related involvement

[betatim]

Can you explain a bit more what you mean with "changing device of an array"? Moving an array from one device to another should be something we can do no? Is this about assigning a difference device to some_array.device?

It probably makes sense to disable/raise exceptions for ways to "move" an array from one device to another in array-api-strict that are likely to be mistakes/unsupported by any real array library. Assigning to .device seems like one such thing.

Is that what you had in mind?

[lucascolley]

1. x.to_device(some_other_device)
2. asarray(x, device=some_other_device)

[lucascolley]

Moving an array from one device to another should be something we can do no?

Probably by default, yes, as I suggested above, but @ogrisel was suggesting that it may be nice to be able to disable that too, to more easily be able to tell if one device is used throughout.

[ogrisel]

I don't recall this particular conversation nor having expressed the need for such an option. In scikit-learn we typically check that the device of the returned arrays match our expectations (often: same device as input). I think this is good enough.

Do you recall a more precise context maybe?

[ogrisel]

Also, the following work without raising an exception in torch via array-api-compat:

>>> import array_api_compat.torch as xp
>>> xp.asarray(xp.ones(shape=3, device="mps"), device="cpu")
tensor([1., 1., 1.])
>>> xp.asarray(xp.ones(shape=3, device="cpu"), device="mps")
tensor([1., 1., 1.], device='mps:0')

[lucascolley]

hmm, feel free to close if it doesn't seem useful! I definitely remember you asking whether array-api-strict already had the capability to detect and report any device transfers.

[ogrisel]

Ah, indeed, now I remember: a way to make it possible to check that array API consuming some code would not cause any potentially costly device transfers under the hood.

Indeed, that could be helpful, however it might be tricky to use properly. For instance, since array API does not specify RNG sampling API, we often generate random numbers using NumPy and then transfer to the namespace of the input data in scikit-learn functions / methods. These kinds of transfers we want to allow because we don't have any alternative. But we might want to check that our code does not silently transfer the whole input data buffer to another device internally, in the same function call. Similarly, we might want to check that the output arrays are directly allocated on the target device rather than being allocated on some intermediate device before transferring to the target device.

[ogrisel]

That might be a nice to have, but it's not a top priority though.

[betatim]

Maybe a nice way to use this would be a context manager that you can use in a test to express "in this code block no device transfers should occur" or even capture information about the transfers that do occur so you can then either assert that none happened or only "allowed ones" (e.g. random number related ones). The inspiration comes from how pytest's warning capture works, I'd see if we can copy that (it is "tried and tested" in terms of UX)

with array_api_strict.fail_on_device_transfer:
  # code that shouldn't perform transfers
  # This would lead t o an exception:
  xp.asarray(xp.ones(shape=3, device="mps"), device="cpu")

with array_api_strict.capture_device_transfers as transfers:
  # This leads to an entry in `transfers`
  xp.asarray(xp.ones(shape=3, device="mps"), device="cpu")

assert not transfers
# or if you want to allow one
assert len(transfers) == 1
assert transfers[0].origin_device == "cpu"
assert transfers[0].target_device == "mps"

Just an idea

[lucascolley]

very nice idea!

[ogrisel]

I think that the title of this issue should be renamed to:
"Add a config manager to record all device transfers" (or raise an exception if any) and link to #168 (comment) from the description.