-
-
Notifications
You must be signed in to change notification settings - Fork 5.5k
ENH: ndimage: add array API standard support #21150
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
ebe7ec4
to
de94be0
Compare
42f6224
to
fec58d7
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @ev-br!
This is a huge PR that looks good at a high level. I only reviewed the array API usage and the new utility wrappers/functions added in _lib/_array_api.py
. Some of those indeed seem useful and make things a lot easier (the testing functions in particular), but some seem unnecessary. It's better to use idiomatic functions where possible, rather than rewrite everything that we had as np.xxx
as xp_xxx
(e.g., xp_imag
, xp_isscalar
). So I'd like to see those removed. The resulting code changes should actually (a) be minor, and (b) make the code more readable.
Thanks Ralf! Have to admit I don't like where this is going, too, and I'm going to rework it all to avoid |
c23d5cd
to
4f80326
Compare
I'm probably missing something very simple. So the point of However IRL, after an update which makes tests pass under
|
xp_test = array_namespace(x) # explain here why we need the wrapped namespace, e.g. # `cupy` needs `concat`
# use `xp_test` instead of `xp` You should ensure that the arrays which are passed into the functions being tested are still created with the raw namespace Sometimes it is better to avoid using this pattern by using NumPy functions and converting to That is for failures arising from the test functions themselves. If there are failures arising from the actual functions being tested, then something is going wrong with coercing to the wrapped namespace. |
Can we please flag what's essential to change and what can be left for a follow-up? At this point I'm no longer clear on what is what. |
Sure, I'll have to take another look through your round of changes first. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks Evgeni, again this is a huge amount of work! I have resolved all addressed comments - it should be clear now what I expect to be changed before merge here. I've opened gh-21280 to track follow-ups.
265f74c
to
5cfe3f0
Compare
OK, 5cfe3f0 addresses the two remaining comments. TBH, I don't necessarily agree that the Let me know if you're happy with the result @lucascolley . The last thing is probably to squash the commits some --- 60+ commits is not useful here, so it's either a squash-merge (5kLOC) or separate into a pair of commits: one to add the delegation layer, the other to overhaul the tests (still 5kLOC). I'm leaning towards the second option, just so that the delegation stuff is more discoverable. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let me know if you're happy with the result @lucascolley . The last thing is probably to squash the commits some --- 60+ commits is not useful here, so it's either a squash-merge (5kLOC) or separate into a pair of commits: one to add the delegation layer, the other to overhaul the tests (still 5kLOC). I'm leaning towards the second option, just so that the delegation stuff is more discoverable.
LGTM! Squashing into a pair of commits sounds sensible to me 👍
The internal ndimage implementation remains numpy-only, and is intact. - On CuPy, delegate ndimage.func to cupyx.scipy.ndimage.func - On other backends, convert inpit arguments to numpy arrays, run the numpy-only internal function, and convert the results back - On JAX, delegate for map-coordinates and convert for the rest. This way, each ndimage function gets its _dispatcher function, which knows which arguments are array-like, and deduces the array namespace based on inputs. This deduction stage needs to be function-specific. The rest is regular, and is done in _support_alternative_backends.py Constructing the public API goes through an indirection layer: - _ndimage_api.py collect imports from private implementation modules; - _support_alternative_backends.py handles decorates implementor functions; - __init__.py import from _support_alternative_backends.py
Since ndimage/_{implementation module} are all numpy-only, they need to make sure their arguments are not cupy/torch/jax/etc arrays. Previously only numpy arrays or python array-convertible structures were possible, so several missing asarray calls were not a problem and conversion was done ad hoc in numpy calls in the 'computational' parts of the functions.
1db983c
to
3bd257a
Compare
... and assert_array_almost equal. This is to reduce the size of the diff when tests are converted to use xp_assert_* infrastructure.
The diff is large but largely mechanical.
Also silence several annoying but harmless warnings: cupy likes to warn that jittify is experimental, mypy is acting up on the import structure etc; there's no need to fail the tests suite on these.
3bd257a
to
8931970
Compare
Rebased. It's five commits in the end: delegation, changes to ndimage itself, _lib testing infra, ndimage/testing refactors, and CI + mypy, all in separate commits. |
Landed - thanks Evgeni! |
Thanks for reviewing this Lucas! |
@mdhaber if you are still interested in looking at generalising the delegation1 infrastructure across here, Footnotes
|
Reference issue
This is on top of #21091, so only the last commit (de94be0) is relevant.What does this implement/fix?
Allow
ndimage
functions consume and return arrays from API compliant namespaces.cupyx.scipy.ndimage
namesakes. See ENH: ndimage: delegate to CuPy #21091 for the initial discussion.Additional information
Since
ndimage
is mainly compiled code with only a very small amount of python (mostly prepare/permute the inputs), we keep the original ndimage code almost intact. Wrapping/unwrapping etc is done in a decorator.