ENH: ndimage: delegate to CuPy #21091

ev-br · 2024-07-01T10:55:23Z

Reference issue

cross-ref #21012, #20772

What does this implement/fix?

This is a POC of the layer of dispatch/delegation from scipy.ndimage to cupyx.scipy.ndimage, for a full scipy submodule. ndimage is probably clearer than signal in that it's all compiled code.

Several things to note:

There are no changes to the ndimage code itself, the diff is tweaking tests and adding a the delegation layer only.
Test changes are mechanical: port the test suite to the xp_assert framework, add xfails for CuPy issues
_support_alternative_backends.py is for now copy-pasted from ENH: array types, signal: delegate to CuPy and JAX for correlations and convolutions #20772; the next step would be to deduplicate with the machinery from ENH: linalg/signal/special: unify approach for array API dispatching #21012; I think some simplifications are possible, too.
One difference from ENH: linalg/signal/special: unify approach for array API dispatching #21012 is that instead of a mapping dictionary here we use a dedicated private module, _dispatchers.py.
In general, the delegation machinery is this:
- _ndimage_api.py collects imports from implementer modules instead of __init__.py
- _supports_alternative_backends.py decorates things from _ndimage_api with the cupy dispatch
- __init__.py imports decorated names from _supports_alternative_backends.py
Delegation is based on _dispatcher functions. The full set is worked out for the whole submodule, so we can assess if we can simplify and possibly avoid dispatching on a full signature (I don't think we can, but am open to pleasant surprises).
I am not hooking it to any CI yet, because we do not test CuPy on CI anyway.

Technical TODOs:

work out the delegation to jax.scipy.ndimage --- unlike cupyx.scipy.ndimage, it only implements map_coordinates?
make it run gracefully under dev.py -b all

Additional information

mdhaber · 2024-07-01T12:13:49Z

scipy/ndimage/tests/test_filters.py

@@ -107,25 +129,25 @@ def _validate_complex(self, array, kernel, type2, mode='reflect', cval=0):
        # test correlate output dtype
        output = correlate(array, kernel, output=type2)
        assert_array_almost_equal(expected, output)
-        assert_equal(output.dtype.type, type2)
+        assert output.dtype.type == type2


IIUC, the dtypes passed in here will always be NumPy dtypes, but this takes advantage of the fact that CuPy dtypes can be compared with NumPy dtypes and this PR is only for CuPy backend?

Yeah. CuPy dtypes in fact are numpy dtypes, so this works for cupy. Won't work for dtypes from array-api-strict though, so if we want to make it work in this PR, the diff is going to grow by another ~500 mechanical changes :-).

scipy/ndimage/tests/test_filters.py

mdhaber · 2024-07-01T12:16:21Z

scipy/ndimage/tests/test_filters.py


        # test correlate with pre-allocated output
-        output = np.zeros_like(array, dtype=type2)
+        output = xp.zeros_like(array, dtype=type2)
        correlate(array, kernel, output=output)
        assert_array_almost_equal(expected, output)


Does assert_array_almost_equal work for CuPy arrays?

No. This is a shim from _lib.

ev-br · 2024-07-03T09:04:18Z

speaking of dtypes: what is the recommended way to to declare a dtype across all array api compatible backends?

The pattern is

@pytest.mark.parametrize(dt, [np.float32, np.float64])
def test_somethin(xp):
    arr = xp.zeros((3, 4), dtype=dt)
    ...

and it fails to understand the dtype under array_api_strict. I tried specifying strings, and do xp.dtype("float64") but that fails, too.

lucascolley · 2024-07-03T09:33:56Z

xp.float64 should work

(However, since we are dealing with the raw namespaces rather than the array-api-compat wrapped versions initially in tests, this is what is blocking us from adding dask, as there is no dask.array.float64. Maybe we will need to invent our own solution for robustly parametrising dtypes, but I'm going to ask in the Dask community meeting if they can help us out)

ev-br · 2024-07-03T09:38:29Z

xp.float64 should work

In the test itself yes, but how to @parametrize: xp is not known outside of the test function.

lucascolley · 2024-07-03T09:42:28Z

Oops 🤦‍♂️ try getattr(xp, dtype) and parametrize with strings like "float64".

ev-br · 2024-07-03T09:49:38Z

No dice:

diff --git a/scipy/ndimage/tests/test_interpolation.py b/scipy/ndimage/tests/test_interpolation.py
index db3479dbb1..63e816d7d8 100644
--- a/scipy/ndimage/tests/test_interpolation.py
+++ b/scipy/ndimage/tests/test_interpolation.py
@@ -194,11 +194,12 @@ class TestNdimageInterpolation:
         assert_array_almost_equal(out, [0, 4, 1, 3])
 
     @pytest.mark.parametrize('order', range(0, 6))
-    @pytest.mark.parametrize('dtype', [np.float64, np.complex128])
+    @pytest.mark.parametrize('dtype', ["float64", "complex128"])
     def test_geometric_transform05(self, order, dtype, xp):
         if is_cupy(xp):
             pytest.xfail("CuPy does not have geometric_transform")
 
+        dtype = getattr(xp, "dtype")
         data = xp.asarray([[1, 1, 1, 1],
                          [1, 1, 1, 1],
                          [1, 1, 1, 1]], dtype=dtype)

produces

$ python dev.py test -t scipy/ndimage/tests/test_interpolation.py -b array_api_strict -- -k 'test_geometric_transform05'
...
_____________________________________________________________ TestNdimageInterpolation.test_geometric_transform05[float64-0-array_api_strict] _____________________________________________________________
scipy/ndimage/tests/test_interpolation.py:202: in test_geometric_transform05
    dtype = getattr(xp, "dtype")
E   AttributeError: module 'array_api_strict' has no attribute 'dtype'. Did you mean: '_dtypes'?
        dtype      = 'float64'

j-bowhay · 2024-07-03T10:10:48Z

diff --git a/scipy/ndimage/tests/test_interpolation.py b/scipy/ndimage/tests/test_interpolation.py
index db3479dbb1..63e816d7d8 100644
--- a/scipy/ndimage/tests/test_interpolation.py
+++ b/scipy/ndimage/tests/test_interpolation.py
@@ -194,11 +194,12 @@ class TestNdimageInterpolation:
         assert_array_almost_equal(out, [0, 4, 1, 3])
 
     @pytest.mark.parametrize('order', range(0, 6))
-    @pytest.mark.parametrize('dtype', [np.float64, np.complex128])
+    @pytest.mark.parametrize('dtype', ["float64", "complex128"])
     def test_geometric_transform05(self, order, dtype, xp):
         if is_cupy(xp):
             pytest.xfail("CuPy does not have geometric_transform")
 
+        dtype = getattr(xp, "dtype")
         data = xp.asarray([[1, 1, 1, 1],
                          [1, 1, 1, 1],
                          [1, 1, 1, 1]], dtype=dtype)

I think there is a type in your patch, should it not be dtype = getattr(xp, dtype) instead of dtype = getattr(xp, "dtype")

ev-br · 2024-07-03T10:12:48Z

Ah, that's it. Thanks Jake!

lucascolley

I've finished one of the big test files - there are a lot of occurrences of indentation being thrown off by adding xp.asarray. A few comments otherwise.

scipy/ndimage/tests/test_interpolation.py

ev-br

Review comments addressed.

scipy/ndimage/tests/test_interpolation.py

scipy/ndimage/tests/test_morphology.py

ev-br · 2024-07-13T14:13:13Z

OK, I hopefully reformatted all instances of misaligned arrays.

lucascolley

getting there! this is a huge diff 😅

scipy/ndimage/tests/test_interpolation.py

scipy/ndimage/tests/test_morphology.py

ev-br · 2024-07-13T20:54:42Z

this is a huge diff

This is all going to be mechanical, they said :-).

ev-br · 2024-07-14T09:54:03Z

All right, all the test formatting etc should be in shape now.

scipy/ndimage/tests/test_interpolation.py

ev-br · 2024-07-14T14:26:08Z

The last round of comments is addressed.

ev-br · 2024-07-15T13:06:04Z

FYI @lucascolley turning other backends on requires further test tweaks. I'll sync tests here when #21150 reaches a natural sync point.

ev-br · 2024-07-16T08:58:45Z

All tests are in sync now. Sorry for the back-and-forth @lucascolley

lucascolley · 2024-07-17T10:08:20Z

could you take a look at the CI failures here? I assume most of it is adding the appropriate skips, to be removed in the next PR.

ev-br · 2024-07-17T13:55:04Z

Ha, this one is a funny one. TL;DR: skip_xp_backends do not add up.

Consider test_white_tophat01 test in test_morphology.py, which fails here. The test module is

from scipy.conftest import array_api_compatible
skip_xp_backends = pytest.mark.skip_xp_backends
pytestmark = [array_api_compatible, pytest.mark.usefixtures("skip_xp_backends"),
              # XXX: only CuPy delegation is implemented
              skip_xp_backends("jax.numpy", "torch", "array_api_strict"),
]

<snip>

   @skip_xp_backends("jax.numpy", reason="output array is read-only.")
   def test_white_tophat01(self, xp):
        .....

The decorator clears other skips, so the test is only skipped on jax --- and fails miserably on torch and array-api-strict.

It seems to me that the attempt to break a large PR into a "stack" of this one and #21150, which works great on some other projects, simply does not work in SciPy. Or at least does not work here anyway.

I'm open to suggestions on how to move this forward. Preferably in a way that won't require me adding and immediately removing ~100 skips.

lucascolley · 2024-07-17T14:17:14Z

Would you like to just consolidate into 1 PR? The diff is so big anyway that making it a bit bigger doesn't really make it any more difficult to review.

RE: multi-level skips - yes, this has caught me out before. I haven't rushed to work around this, though, as I think the mental model of 'the source of truth for skips is the closest decorator to the function's level, if one exists' is okay to grasp, once you've been told about it. It just isn't the most intuitive model.

ev-br · 2024-07-17T14:32:20Z

OK, let's continue in #21150 indeed.

The bulk of the diff is just the same:

test tweaks are the same
the decoration/delegation story is mostly the same, the main difference is that the full version just needs a small additional dance to take care of various non-array return values.

'the source of truth for skips is the closest decorator to the function's level, if one exists' is okay to grasp, once you've been told about it. It just isn't the most intuitive model.

The main footgun IME is that repeated skips tend to be fall into cracks: if you skip X in pytestmark, need to skip Y for a particular test, you've got to skip both X and Y --- and then when you undo the pytestmark, the test keeps skipping both.

ev-br requested review from larsoner and rgommers as code owners July 1, 2024 10:55

github-actions bot added scipy.ndimage scipy._lib 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 1, 2024

ev-br requested review from lucascolley and mdhaber and removed request for rgommers, larsoner and lucascolley July 1, 2024 10:55

This was referenced Jul 1, 2024

ENH: linalg/signal/special: unify approach for array API dispatching #21012

Closed

ENH: array types, signal: delegate to CuPy and JAX for correlations and convolutions #20772

Merged

lucascolley changed the title ~~ndimage cupy delegation~~ ENH/POC: ndimage: delegate to CuPy Jul 1, 2024

lucascolley added the enhancement A new feature or improvement label Jul 1, 2024

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch from 21290a1 to eaa9a0a Compare July 1, 2024 11:19

mdhaber reviewed Jul 1, 2024

View reviewed changes

scipy/ndimage/tests/test_filters.py Outdated Show resolved Hide resolved

mdhaber reviewed Jul 1, 2024

View reviewed changes

ev-br mentioned this pull request Jul 1, 2024

RFC: require alternative backends to pass SciPy's tests when "dispatching"? #21031

Closed

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch from 07c51a0 to fd28566 Compare July 3, 2024 09:28

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch from 0761236 to 9b013f4 Compare July 3, 2024 13:22

lucascolley requested changes Jul 12, 2024

View reviewed changes

lucascolley removed request for andyfaff and ilayn July 12, 2024 12:32

ev-br commented Jul 13, 2024

View reviewed changes

scipy/ndimage/tests/test_interpolation.py Show resolved Hide resolved

lucascolley requested changes Jul 13, 2024

View reviewed changes

TST: ndimage: address review comments

2a6944f

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch from 4fd7284 to 2a6944f Compare July 13, 2024 14:12

lucascolley requested changes Jul 13, 2024

View reviewed changes

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch from 1bae880 to 021addf Compare July 14, 2024 09:52

lucascolley requested changes Jul 14, 2024

View reviewed changes

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch 2 times, most recently from 44406fd to 705d9c9 Compare July 14, 2024 14:25

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch 2 times, most recently from 05c5e55 to 37655f9 Compare July 14, 2024 15:06

TST: ndimage: sync test files for array API

f4e5d46

ev-br force-pushed the ndimage_array_api_tests_and_dispatch branch from 37655f9 to f4e5d46 Compare July 16, 2024 08:57

ev-br closed this Jul 17, 2024

lucascolley mentioned this pull request Jul 28, 2024

TST/MAINT: ndimage: follow-ups for array API standard support #21280

Open

5 tasks

ev-br mentioned this pull request Sep 4, 2024

MAINT: ndimage: rename dispatch -> delegate #21506

Merged

ev-br removed this from the 1.15.0 milestone Oct 25, 2024

Uh oh!

ENH: ndimage: delegate to CuPy #21091

ENH: ndimage: delegate to CuPy #21091

Uh oh!

Conversation

ev-br commented Jul 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reference issue

What does this implement/fix?

Additional information

Uh oh!

mdhaber Jul 1, 2024

Choose a reason for hiding this comment

Uh oh!

ev-br Jul 2, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mdhaber Jul 1, 2024

Choose a reason for hiding this comment

Uh oh!

ev-br Jul 1, 2024

Choose a reason for hiding this comment

Uh oh!

ev-br commented Jul 3, 2024

Uh oh!

lucascolley commented Jul 3, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ev-br commented Jul 3, 2024

Uh oh!

lucascolley commented Jul 3, 2024

Uh oh!

ev-br commented Jul 3, 2024

Uh oh!

j-bowhay commented Jul 3, 2024

Uh oh!

ev-br commented Jul 3, 2024

Uh oh!

lucascolley left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ev-br left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ev-br commented Jul 13, 2024

Uh oh!

lucascolley left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ev-br commented Jul 13, 2024

Uh oh!

ev-br commented Jul 14, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ev-br commented Jul 1, 2024 •

edited

Loading

lucascolley commented Jul 3, 2024 •

edited

Loading

lucascolley left a comment •

edited

Loading

lucascolley commented Jul 17, 2024 •

edited

Loading