Merge pull request #25536 from mhvk/fft-as-gufunc

MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument
numpy · Jan 19, 2024 · 539dafa · 539dafa
2 parents 72879d3 + 570f272
commit 539dafa
Show file tree

Hide file tree

Showing 8 changed files with 739 additions and 272 deletions.
diff --git a/numpy/fft/_pocketfft.py b/numpy/fft/_pocketfft.py
diff --git a/numpy/fft/_pocketfft.pyi b/numpy/fft/_pocketfft.pyi
@@ -13,27 +13,31 @@ def fft(
     n: None | int = ...,
     axis: int = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def ifft(
     a: ArrayLike,
     n: None | int = ...,
     axis: int = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def rfft(
     a: ArrayLike,
     n: None | int = ...,
     axis: int = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def irfft(
     a: ArrayLike,
     n: None | int = ...,
     axis: int = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[float64] = ...,
 ) -> NDArray[float64]: ...
 
 # Input array must be compatible with `np.conjugate`
@@ -42,67 +46,77 @@ def hfft(
     n: None | int = ...,
     axis: int = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[float64] = ...,
 ) -> NDArray[float64]: ...
 
 def ihfft(
     a: ArrayLike,
     n: None | int = ...,
     axis: int = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def fftn(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def ifftn(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def rfftn(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def irfftn(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[float64] = ...,
 ) -> NDArray[float64]: ...
 
 def fft2(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def ifft2(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def rfft2(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[complex128] = ...,
 ) -> NDArray[complex128]: ...
 
 def irfft2(
     a: ArrayLike,
     s: None | Sequence[int] = ...,
     axes: None | Sequence[int] = ...,
     norm: _NormKind = ...,
+    out: None | NDArray[float64] = ...,
 ) -> NDArray[float64]: ...