CuPy support #6
CuPy support #6
Conversation
I still need to factor out uses of 'np' in other places in the code. I'm also not completely sure about the solution for asarray(). asarray() can't just infer the namespace from the input like every other function can. For now, I have added a namespace keyword argument to it, which defaults to numpy.
| try: | ||
| xp = get_namespace(obj, _use_compat=False) | ||
| except ValueError: | ||
| # TODO: What about lists of arrays? |
There was a problem hiding this comment.
What you have seems reasonable. Libraries should be passing namespace=numpy if they want this to work, or otherwise handle lists/sequences themselves.
|
I'm looking at |
|
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. |
There will also be a 'cupy' submodule. Common code will be factored out into the 'common' submodule.
Full support still needs to be tested, and also we need to double check if any of the aliases are unnecessary for cupy.
Now instead of guessing the array library from the input with get_namespace, the array library is hard-coded into the wrapped function based on which subnamespace it is imported from.
|
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. |
There was a problem hiding this comment.
I like the approach Aaron, it's pretty easy to understand and it should be extensible.
The one thing that comes to mind as a possible concern is that the get_xp decorator may have some overhead that's larger than desired - if so that could be fairly easily changed to a codegen approach where xp gets replaced by np or cp it looks like.
README.md
Outdated
|
|
||
| Each will include all the functions from the normal NumPy/CuPy namespace, | ||
| except that functions that are part of the array API are wrapped so that they | ||
| have the correct array API behavior. In each case, the array object |
There was a problem hiding this comment.
incomplete sentence here - and it seems to start saying something that's fairly important.
|
|
||
| import cupy as cp | ||
|
|
||
| # TODO: Should we reject ndarray subclasses? |
There was a problem hiding this comment.
I am not aware of "bad" cupy.ndarray subclasses that need to be avoided.
|
|
||
| import numpy as np | ||
|
|
||
| # TODO: Should we reject ndarray subclasses? |
There was a problem hiding this comment.
Always a question mark. With "reject", do you mean "convert to ndarray", or "raise", or "ignore", or ....?
There was a problem hiding this comment.
My thinking here was specifically for get_namespace, although this helper is also used in asarray. Right now it does accept them, because it uses isinstance. I've no idea if an ndarray subclass could implement behavior that would break things here. Presumably if they are properly Liskov substitutable, they shouldn't. My guess is that at worst some functions might return ndarray instead of the appropriate subclass, because we don't properly account for subclases when creating new array objects.
But I haven't tested it. It's possible some people one of the units libraries (I forget which ones subclass ndarray) might care about this at some point, though.
There was a problem hiding this comment.
Presumably if they are properly Liskov substitutable, they shouldn't.
np.matrix for one isn't, that's a major issue. IIRC also np.ma.MaskedArray isn't.
|
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? |
|
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:) |
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.