Description
Motivation: in scikit-learn, we run array API compliance tests both with array_api_strict
and array_api_compat.torch
. For the latter we run tests both with cuda
, mps
and cpu
devices.
Testing with torch is important because it helps us reveal occurrences of problems related to device handling. For example let's consider the following problematic function:
def sublinear_mask(data):
return data <= xp.linspace(data[0], data[-1], num=data.shape[0])
Calling this with:
sublinear_mask(xp.asarray([0, 1, 2, 2, 5, 5], device="mps"))
raises a RuntimeError
because PyTorch does not implicitly move data across devices. Hence the bug should be fixed by changing the function to:
def sublinear_mask(data):
return data <= xp.linspace(data[0], data[-1], num=data.shape[0], device=data.device)
However, not all scikit-learn contributors have access to a machine with a non-cpu device (e.g. "mps" or "cuda") and therefore they have no easy way to detect this family of bugs by running their tests on their local laptop and only discover issues on the CI and need to use tools such as google colab to debug and refine their code instead of their regular dev environment.
To reduce this friction, it would be nice if array-api-strict
could accept to create arrays with a = xp.asarray([1, 2, 3], device="virtual_device_a")
and b = xp.asarray([1, 2, 3], device="virtual_device_b")
and would raise RuntimeError
on operations that combine arrays from different devices as pytorch does.