data-apis / array-api-compat

Compatibility layer for common array libraries to support the Array API
https://data-apis.org/array-api-compat/
MIT License
76 stars 26 forks source link

CuPy support #6

Closed asmeurer closed 1 year ago

asmeurer commented 1 year ago

This refactors the code to be array library independent, so that it can work with either NumPy or CuPy.

Other libraries could be supported as well, although the wrapping for each library would be specific to the differences between that library and the array API.

Fixes #3.

rgommers commented 1 year ago

I'm looking at get_xp and thinking there's still a piece of the puzzle missing here. Should I wait until that's in place, or am I just missing how it works?

asmeurer commented 1 year ago

I haven't added any cupy functionality yet if that's what you mean. I also need to make sure that NumPy is only an optional dependency of all the code here.

asmeurer commented 1 year ago

This is ready to go. Anyone want to review the overall design here? I've described it in the README.

I still need to set up some CI here, but I've tested the NumPy and CuPy support with the array API test suite and the only issues are things which can't really be fixed by wrapping.

asmeurer commented 1 year ago

Here's an idea of the overhead of get_xp:

In [1]: import numpy as np

In [2]: import array_api_compat.numpy as npc

In [3]: a = np.asarray(1.)

In [4]: %timeit np.arccos(a)
435 ns ± 8.46 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

In [5]: %timeit npc.acos(a)
737 ns ± 4.31 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

It looks like it's about 300 ns per function call. Does this seem acceptable or should we investigate faster approaches?

rgommers commented 1 year ago

Thanks @asmeurer. I think that that is fine for now; in case it becomes a problem because this function is used a lot, I think there is a fairly straightforward way to improve it. Let's first get this used, once a user complains about the overhead I guess we'll have made very nice progress with adoption:)