From c05d7fc4c26e05416767e6ae0c6fac3a4b6f9663 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Mon, 20 Jan 2025 02:56:23 -0800 Subject: [PATCH 1/2] feat: add `dtype` kwarg support to `fftfreq` and `rfftfreq` Closes: https://github.com/data-apis/array-api/issues/717 --- src/array_api_stubs/_draft/fft.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/src/array_api_stubs/_draft/fft.py b/src/array_api_stubs/_draft/fft.py index 7a4538ccb..1c8ed1358 100644 --- a/src/array_api_stubs/_draft/fft.py +++ b/src/array_api_stubs/_draft/fft.py @@ -551,7 +551,14 @@ def ihfft( """ -def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> array: +def fftfreq( + n: int, + /, + *, + d: float = 1.0, + dtype: Optional[dtype] = None, + device: Optional[device] = None, +) -> array: """ Computes the discrete Fourier transform sample frequencies. @@ -568,13 +575,15 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> ar window length. d: float sample spacing between individual samples of the Fourier transform input. Default: ``1.0``. + dtype: Optional[dtype] + output array data type. Must be a real-valued floating-point data type. If ``dtype`` is ``None``, the output array data type must be the default real-valued floating-point data type. Default: ``None``. device: Optional[device] device on which to place the created array. Default: ``None``. Returns ------- out: array - an array of shape ``(n,)`` containing the sample frequencies. The returned array must have the default real-valued floating-point data type. + an array of shape ``(n,)`` containing the sample frequencies. Notes ----- @@ -586,7 +595,14 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> ar """ -def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> array: +def rfftfreq( + n: int, + /, + *, + d: float = 1.0, + dtype: Optional[dtype] = None, + device: Optional[device] = None, +) -> array: """ Computes the discrete Fourier transform sample frequencies (for ``rfft`` and ``irfft``). @@ -605,13 +621,15 @@ def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> a window length. d: float sample spacing between individual samples of the Fourier transform input. Default: ``1.0``. + dtype: Optional[dtype] + output array data type. Must be a real-valued floating-point data type. If ``dtype`` is ``None``, the output array data type must be the default real-valued floating-point data type. Default: ``None``. device: Optional[device] device on which to place the created array. Default: ``None``. Returns ------- out: array - an array of shape ``(n//2+1,)`` containing the sample frequencies. The returned array must have the default real-valued floating-point data type. + an array of shape ``(n//2+1,)`` containing the sample frequencies. Notes ----- From 0a76b141b4891e3ede89e00d9e6ce40e4b2434f5 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Mon, 20 Jan 2025 03:02:36 -0800 Subject: [PATCH 2/2] fix: add missing type --- src/array_api_stubs/_draft/fft.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/array_api_stubs/_draft/fft.py b/src/array_api_stubs/_draft/fft.py index 1c8ed1358..3d1f4d9c7 100644 --- a/src/array_api_stubs/_draft/fft.py +++ b/src/array_api_stubs/_draft/fft.py @@ -15,7 +15,7 @@ "ifftshift", ] -from ._types import Tuple, Union, Sequence, array, Optional, Literal, device +from ._types import Tuple, Union, Sequence, array, Optional, Literal, dtype, device def fft(