Skip to content

Conversation

ev-br
Copy link
Member

@ev-br ev-br commented Jul 10, 2024

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.

  • on cupy, delegate to cupyx.scipy.ndimage namesakes. See ENH: ndimage: delegate to CuPy #21091 for the initial discussion.
  • on other backends, cast inputs to numpy, do computations and cast the output back.

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.

@github-actions github-actions bot added scipy.signal scipy.ndimage scipy._lib CI Items related to the CI tools such as CircleCI, GitHub Actions or Azure Meson Items related to the introduction of Meson as the new build system for SciPy array types Items related to array API support and input array validation (see gh-18286) labels Jul 10, 2024
@ev-br ev-br force-pushed the ndimage_array_api_inline branch from ebe7ec4 to de94be0 Compare July 10, 2024 14:45
@lucascolley lucascolley changed the title WIP: ndimage array api support ENH: ndimage: add array API standard support Jul 10, 2024
@lucascolley lucascolley marked this pull request as draft July 10, 2024 15:09
@lucascolley lucascolley added enhancement A new feature or improvement and removed CI Items related to the CI tools such as CircleCI, GitHub Actions or Azure Meson Items related to the introduction of Meson as the new build system for SciPy scipy.signal labels Jul 10, 2024
@ev-br ev-br force-pushed the ndimage_array_api_inline branch 2 times, most recently from 42f6224 to fec58d7 Compare July 11, 2024 11:01
Copy link
Member

@rgommers rgommers left a 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.

@ev-br
Copy link
Member Author

ev-br commented Jul 12, 2024

I only reviewed the array API usage and the new utility wrappers/functions added in _lib/_array_api.py

Thanks Ralf! Have to admit I don't like where this is going, too, and I'm going to rework it all to avoid xp_* additions at an expense of some more meta-stuff.

@ev-br ev-br mentioned this pull request Jul 13, 2024
2 tasks
@ev-br ev-br force-pushed the ndimage_array_api_inline branch 4 times, most recently from c23d5cd to 4f80326 Compare July 14, 2024 10:06
@ev-br
Copy link
Member Author

ev-br commented Jul 14, 2024

I'm probably missing something very simple. So the point of array_api_strict is to provide a lowest common denominator for all array API libraries, and if tests pass with -b array_api_strict, they should pass under all backends, correct?

However IRL, after an update which makes tests pass under array_api_strict == 2, other backends fail left right and center with missing array API attributes and methods:

$ python dev.py test -s ndimage -b cupy -j 12
...
FAILED scipy/ndimage/tests/test_measurements.py::test_sum10[cupy] - AttributeError: module 'cupy' has no attribute 'bool'
FAILED scipy/ndimage/tests/test_interpolation.py::TestNdimageInterpolation::test_affine_transform26[2-cupy] - AttributeError: module 'cupy' has no attribute 'concat'
...

Am I missing something simple?

@lucascolley
Copy link
Member

lucascolley commented Jul 14, 2024

However IRL, after an update which makes tests pass under array_api_strict == 2, other backends fail left right and center with missing array API attributes and methods:

array_api_compatible parametrises tests with the raw (i.e. un-array_api_compat-wrapped) namespaces. If you need to use functions which exist in the standard but not in the main namespaces of the libraries we test with, you can use the pattern:

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 xp rather than xp_test.

Sometimes it is better to avoid using this pattern by using NumPy functions and converting to xp at a later point.


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.

@ev-br
Copy link
Member Author

ev-br commented Jul 28, 2024

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.

@lucascolley
Copy link
Member

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.

@lucascolley lucascolley added the needs-release-note a maintainer should add a release note written by a reviewer/author to the wiki label Jul 28, 2024
@lucascolley lucascolley added this to the 1.15.0 milestone Jul 28, 2024
Copy link
Member

@lucascolley lucascolley left a 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.

@ev-br ev-br force-pushed the ndimage_array_api_inline branch from 265f74c to 5cfe3f0 Compare July 28, 2024 16:56
@ev-br
Copy link
Member Author

ev-br commented Jul 28, 2024

OK, 5cfe3f0 addresses the two remaining comments.

TBH, I don't necessarily agree that the sumsq change is that useful, on the grounds that it adds essentially dead code. But this is not essential enough for me to argue, and it's changed in the above commit.

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.

Copy link
Member

@lucascolley lucascolley left a 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 👍

ev-br added 2 commits July 29, 2024 09:29
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.
@ev-br ev-br force-pushed the ndimage_array_api_inline branch from 1db983c to 3bd257a Compare July 29, 2024 06:57
ev-br added 3 commits July 29, 2024 09:57
... and assert_array_almost equal.

This is to reduce the size of the diff when tests are converted to use
xp_assert_* infrastructure.
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.
@ev-br ev-br force-pushed the ndimage_array_api_inline branch from 3bd257a to 8931970 Compare July 29, 2024 06:57
@ev-br
Copy link
Member Author

ev-br commented Jul 29, 2024

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.

@lucascolley lucascolley merged commit b2c5b76 into scipy:main Jul 29, 2024
39 checks passed
@lucascolley
Copy link
Member

Landed - thanks Evgeni!

@ev-br
Copy link
Member Author

ev-br commented Jul 29, 2024

Thanks for reviewing this Lucas!

@lucascolley
Copy link
Member

@mdhaber if you are still interested in looking at generalising the delegation1 infrastructure across here, special and beyond, that would be fantastic!

Footnotes

  1. we decided in the community meeting that avoiding the word "dispatching" to refer to this is probably best, now that work on dispatching systems seems to be picking back up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
array types Items related to array API support and input array validation (see gh-18286) enhancement A new feature or improvement scipy._lib scipy.ndimage
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants