Skip to content

Add virtual devices to make it easier for array API consumer to check that they use device correctly #56

Closed
@ogrisel

Description

@ogrisel

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions