Improvements to device support

data-apis / array-api-strict

Strict implementation of the Python array API (previously numpy.array_api)

http://data-apis.org/array-api-strict/

Other

11 stars 5 forks source link

Improvements to device support #70

Open asmeurer opened 1 month ago

asmeurer commented 1 month ago

https://github.com/data-apis/array-api-strict/pull/59 added support for basic devices. Some improvements that could be made:

Factor out the device checking logic (into a decorator?)
Add devices that don't support certain dtypes (non-support for float64/complex128 being the most common)
Allow the user to manually create custom "devices" with certain properties (probably via the flags API).
Better testing for device support. This ideally should go in array-api-tests, but currently it doesn't test devices at all, so it would be useful to have some basic tests here.

ogrisel commented 1 month ago

I tried to run the scikit-learn tests with array-api-strict on the non-default device and I got some failures that do not happen with PyTorch and non-default device. So I suspect there is something fishy happening with the new devices in array-api-strict but I have not investigated the root cause yet.

https://github.com/scikit-learn/scikit-learn/pull/30090

betatim commented 1 month ago

I'll take a look/investigate that failure further.