Repository: dask/dask
Pull request #10665 (2023-12-04 10:54:38)

Fix `dask.array.fft` mismatch with Numpy's interface (add support for `norm` argument)
- Closes  #3280 where `norm` argument was needed to match numpy.fft API. 
- Adds tests to validate that the different values for `norm` match NumPy's.
- Passes `pre-commit run --all-files`


GPUtester: Can one of the admins verify this patch?

Admins can comment `ok to test` to allow this one PR to run or `add to allowlist` to allow all future PRs from the same author to run.
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     15 files  ±0       15 suites  ±0   3h 49m 58s :stopwatch: + 3m 5s
 12 941 tests ±0   11 620 :white_check_mark: + 23     917 :zzz:  -  12  404 :x:  - 10 
160 329 runs  ±0  143 168 :white_check_mark: +138  16 757 :zzz:  - 127  404 :x:  - 10 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/20564414678).

Results for commit a985de36. ± Comparison against base commit cdcddb68.

<details>
  <summary>This pull request <b>removes</b> 28 and <b>adds</b> 28 tests. <i>Note that renamed tests count towards both.</i></summary>

```
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-fft2-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-fftn-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-ifft2-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-ifftn-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-irfft2-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-irfftn-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-rfft2-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-rfftn-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float64-fft2-scipy.fftpack]
dask.array.tests.test_fft ‑ test_wrap_fftns[float64-fftn-scipy.fftpack]
…
```

```
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-fft2-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-fftn-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-ifft2-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-ifftn-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-irfft2-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-irfftn-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-rfft2-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float32-rfftn-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float64-fft2-scipy.fft]
dask.array.tests.test_fft ‑ test_wrap_fftns[float64-fftn-scipy.fft]
…
```
</details>

[test-results]:data:application/gzip;base64,H4sIAP+ap2UC/1XMyw7CIBCF4VdpWLvocBnAlzGUS0LsxVBYGd9dWi3V5f9N5jxJiKNfybUDcenIWmI+y5VkclzmrZnUqlK95v1ONYejb2uxdkNA2p94j4+KGmSjYOJYife8kU9pSdW2v1TmfRt7RvW32zZngKrhZxtQCnnY3/guv9t2maaYaxCjlXCeoeWCC4XU2TAED8oGBcExRDOg1ZKS1xtN5NitHQEAAA==

joanrue: Hi! I'm reaching out to see if there's anything more I can do to facilitate the review process.
joanrue: Hi ! The only failing test is not related to these changes. Can we go ahead with the PR?
lucascolley: this would be useful for SciPy! x-ref https://github.com/data-apis/array-api-compat/pull/139
lucascolley: @phofl @joanrue if we want to maintain support for `scipy.fftpack` without the `norm` argument, how about we:

- Create a new kwarg `allow_fftpack` for `fft_wrap`.
- Start raising a warning instead of an error here for `scipy.fftpack` input, noting that the NumPy API is not matched for the `norm` keyword.
  - Direct the user to use `scipy.fft` instead.
  - If the user still wants to use `scipy.fftpack`, make them pass `allow_fftpack` to silence the warning.
- Handle the failing `norm` with `scipy.fftpack` via `try ... except` or an explicit guard on the `fft_func.__module__` by default.
- Ignore `norm` when `allow_fftpack` is `True`
phofl: Yeah this sounds like a good path forward. I am fine with removing the usage of the legacy module at some point, but we should definitely deprecate this first and not just raise if scipy isn't raising warning either.
joanrue: @lucascolley @phofl thanks for your feedback!

Patch:
diff --git a/dask/array/fft.py b/dask/array/fft.py
index acbb2f926..07aed2ecc 100644
--- a/dask/array/fft.py
+++ b/dask/array/fft.py
@@ -115,7 +115,7 @@ _out_chunk_fns = {
 }
 
 
-def fft_wrap(fft_func, kind=None, dtype=None):
+def fft_wrap(fft_func, kind=None, dtype=None, allow_fftpack=False):
     """Wrap 1D, 2D, and ND real and complex FFT functions
 
     Takes a function that behaves like ``numpy.fft`` functions and
@@ -146,10 +146,18 @@ def fft_wrap(fft_func, kind=None, dtype=None):
     >>> parallel_ifft = dff.fft_wrap(np.fft.ifft)
     """
     if scipy is not None:
-        if fft_func is scipy.fftpack.rfft:
-            raise ValueError("SciPy's `rfft` doesn't match the NumPy API.")
-        elif fft_func is scipy.fftpack.irfft:
-            raise ValueError("SciPy's `irfft` doesn't match the NumPy API.")
+        if fft_func.__module__.startswith("scipy.fftpack"):
+            if not allow_fftpack:
+                warnings.warn(
+                    f"Function {fft_func.__name__} from `scipy.fftpack` does not "
+                    "match NumPy's API and is considered legacy. Please use "
+                    "`scipy.fft` instead. To suppress this warning and allow usage"
+                    ", set `allow_fftpack=True`. Support for `scipy.fftpack` will "
+                    "be deprecated in future releases.",
+                    FutureWarning,
+                )
+            # If allow_fftpack is True, we proceed but we skip passing the norm
+            # argument.
 
     if kind is None:
         kind = fft_func.__name__
@@ -158,7 +166,7 @@ def fft_wrap(fft_func, kind=None, dtype=None):
     except KeyError:
         raise ValueError("Given unknown `kind` %s." % kind)
 
-    def func(a, s=None, axes=None):
+    def func(a, s=None, axes=None, norm=None):
         a = asarray(a)
         if axes is None:
             if kind.endswith("2"):
@@ -185,7 +193,7 @@ def fft_wrap(fft_func, kind=None, dtype=None):
         if _dtype is None:
             sample = np.ones(a.ndim * (8,), dtype=a.dtype)
             try:
-                _dtype = fft_func(sample, axes=axes).dtype
+                _dtype = fft_func(sample, axes=axes, norm=norm).dtype
             except TypeError:
                 _dtype = fft_func(sample).dtype
 
@@ -195,18 +203,18 @@ def fft_wrap(fft_func, kind=None, dtype=None):
 
         chunks = out_chunk_fn(a, s, axes)
 
-        args = (s, axes)
+        args = (s, axes, norm)
         if kind.endswith("fft"):
             axis = None if axes is None else axes[0]
             n = None if s is None else s[0]
-            args = (n, axis)
+            args = (n, axis, norm)
 
         return a.map_blocks(fft_func, *args, dtype=_dtype, chunks=chunks)
 
     if kind.endswith("fft"):
         _func = func
 
-        def func(a, n=None, axis=None):
+        def func(a, n=None, axis=None, norm=None):
             s = None
             if n is not None:
                 s = (n,)
@@ -215,7 +223,7 @@ def fft_wrap(fft_func, kind=None, dtype=None):
             if axis is not None:
                 axes = (axis,)
 
-            return _func(a, s, axes)
+            return _func(a, s, axes, norm)
 
     func_mod = inspect.getmodule(fft_func)
     func_name = fft_func.__name__
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
index cceb40634..7f45fabcf 100644
--- a/dask/array/rechunk.py
+++ b/dask/array/rechunk.py
@@ -375,7 +375,13 @@ def rechunk(
     elif method == "p2p":
         from distributed.shuffle import rechunk_p2p
 
-        return rechunk_p2p(x, chunks)
+        return rechunk_p2p(
+            x,
+            chunks,
+            threshold=threshold,
+            block_size_limit=block_size_limit,
+            balance=balance,
+        )
 
     else:
         raise NotImplementedError(f"Unknown rechunking method '{method}'")


Test patch:
diff --git a/dask/array/tests/test_fft.py b/dask/array/tests/test_fft.py
index 6b07b16c1..4dcd2409a 100644
--- a/dask/array/tests/test_fft.py
+++ b/dask/array/tests/test_fft.py
@@ -81,10 +81,22 @@ def test_fft_n_kwarg(funcname):
 
     assert_eq(da_fft(darr, 5), np_fft(nparr, 5))
     assert_eq(da_fft(darr, 13), np_fft(nparr, 13))
+    assert_eq(da_fft(darr, 13, norm="backward"), np_fft(nparr, 13, norm="backward"))
+    assert_eq(da_fft(darr, 13, norm="ortho"), np_fft(nparr, 13, norm="ortho"))
+    assert_eq(da_fft(darr, 13, norm="forward"), np_fft(nparr, 13, norm="forward"))
     assert_eq(da_fft(darr2, axis=0), np_fft(nparr, axis=0))
     assert_eq(da_fft(darr2, 5, axis=0), np_fft(nparr, 5, axis=0))
-    assert_eq(da_fft(darr2, 13, axis=0), np_fft(nparr, 13, axis=0))
-    assert_eq(da_fft(darr2, 12, axis=0), np_fft(nparr, 12, axis=0))
+    assert_eq(
+        da_fft(darr2, 13, axis=0, norm="backward"),
+        np_fft(nparr, 13, axis=0, norm="backward"),
+    )
+    assert_eq(
+        da_fft(darr2, 12, axis=0, norm="ortho"), np_fft(nparr, 12, axis=0, norm="ortho")
+    )
+    assert_eq(
+        da_fft(darr2, 12, axis=0, norm="forward"),
+        np_fft(nparr, 12, axis=0, norm="forward"),
+    )
 
 
 @pytest.mark.parametrize("funcname", all_1d_funcnames)
@@ -129,7 +141,7 @@ def test_nd_ffts_axes(funcname, dtype):
                 assert_eq(r, er)
 
 
-@pytest.mark.parametrize("modname", ["numpy.fft", "scipy.fftpack"])
+@pytest.mark.parametrize("modname", ["numpy.fft", "scipy.fft"])
 @pytest.mark.parametrize("funcname", all_1d_funcnames)
 @pytest.mark.parametrize("dtype", ["float32", "float64"])
 def test_wrap_ffts(modname, funcname, dtype):
@@ -143,28 +155,24 @@ def test_wrap_ffts(modname, funcname, dtype):
     darr2c = darr2.astype(dtype)
     nparrc = nparr.astype(dtype)
 
-    if modname == "scipy.fftpack" and "rfft" in funcname:
-        with pytest.raises(ValueError):
-            fft_wrap(func)
-    else:
-        wfunc = fft_wrap(func)
-        assert wfunc(darrc).dtype == func(nparrc).dtype
-        assert wfunc(darrc).shape == func(nparrc).shape
-        assert_eq(wfunc(darrc), func(nparrc))
-        assert_eq(wfunc(darrc, axis=1), func(nparrc, axis=1))
-        assert_eq(wfunc(darr2c, axis=0), func(nparrc, axis=0))
-        assert_eq(wfunc(darrc, n=len(darrc) - 1), func(nparrc, n=len(darrc) - 1))
-        assert_eq(
-            wfunc(darrc, axis=1, n=darrc.shape[1] - 1),
-            func(nparrc, n=darrc.shape[1] - 1),
-        )
-        assert_eq(
-            wfunc(darr2c, axis=0, n=darr2c.shape[0] - 1),
-            func(nparrc, axis=0, n=darr2c.shape[0] - 1),
-        )
+    wfunc = fft_wrap(func)
+    assert wfunc(darrc).dtype == func(nparrc).dtype
+    assert wfunc(darrc).shape == func(nparrc).shape
+    assert_eq(wfunc(darrc), func(nparrc))
+    assert_eq(wfunc(darrc, axis=1), func(nparrc, axis=1))
+    assert_eq(wfunc(darr2c, axis=0), func(nparrc, axis=0))
+    assert_eq(wfunc(darrc, n=len(darrc) - 1), func(nparrc, n=len(darrc) - 1))
+    assert_eq(
+        wfunc(darrc, axis=1, n=darrc.shape[1] - 1),
+        func(nparrc, n=darrc.shape[1] - 1),
+    )
+    assert_eq(
+        wfunc(darr2c, axis=0, n=darr2c.shape[0] - 1),
+        func(nparrc, axis=0, n=darr2c.shape[0] - 1),
+    )
 
 
-@pytest.mark.parametrize("modname", ["numpy.fft", "scipy.fftpack"])
+@pytest.mark.parametrize("modname", ["numpy.fft", "scipy.fft"])
 @pytest.mark.parametrize("funcname", all_nd_funcnames)
 @pytest.mark.parametrize("dtype", ["float32", "float64"])
 def test_wrap_fftns(modname, funcname, dtype):
@@ -269,3 +277,17 @@ def test_fftshift_identity(funcname1, funcname2, shape, chunks, axes):
             assert len(each_d_r_chunks) != 1
 
     assert_eq(d_r, d)
+
+
+@pytest.mark.parametrize("modname", ["numpy.fft", "scipy.fft", "scipy.fftpack"])
+def test_scipy_fftpack_future_warning(modname):
+    fft_mod = pytest.importorskip(modname)
+
+    if modname == "scipy.fftpack":
+        # Check that a FutureWarning is raised when using scipy.fftpack with allow_fftpack=False
+        with pytest.warns(
+            FutureWarning, match="does not match NumPy's API and is considered legacy"
+        ):
+            da.fft.fft_wrap(fft_mod.fft, allow_fftpack=False)(np.random.random(16))
+    else:
+        da.fft.fft_wrap(fft_mod.fft, allow_fftpack=False)(np.random.random(16))