Skip to content

Commit

Permalink
feat: add dtype kwarg support to fftfreq and rfftfreq
Browse files Browse the repository at this point in the history
  • Loading branch information
kgryte committed Jan 20, 2025
1 parent 52bbfeb commit c05d7fc
Showing 1 changed file with 22 additions and 4 deletions.
26 changes: 22 additions & 4 deletions src/array_api_stubs/_draft/fft.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
-----
Expand All @@ -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``).
Expand All @@ -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
-----
Expand Down

0 comments on commit c05d7fc

Please sign in to comment.