Skip to content

Commit c05d7fc

Browse files
committed
feat: add dtype kwarg support to fftfreq and rfftfreq
Closes: data-apis#717
1 parent 52bbfeb commit c05d7fc

File tree

1 file changed

+22
-4
lines changed
  • src/array_api_stubs/_draft

1 file changed

+22
-4
lines changed

src/array_api_stubs/_draft/fft.py

+22-4
Original file line numberDiff line numberDiff line change
@@ -551,7 +551,14 @@ def ihfft(
551551
"""
552552

553553

554-
def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> array:
554+
def fftfreq(
555+
n: int,
556+
/,
557+
*,
558+
d: float = 1.0,
559+
dtype: Optional[dtype] = None,
560+
device: Optional[device] = None,
561+
) -> array:
555562
"""
556563
Computes the discrete Fourier transform sample frequencies.
557564
@@ -568,13 +575,15 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> ar
568575
window length.
569576
d: float
570577
sample spacing between individual samples of the Fourier transform input. Default: ``1.0``.
578+
dtype: Optional[dtype]
579+
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``.
571580
device: Optional[device]
572581
device on which to place the created array. Default: ``None``.
573582
574583
Returns
575584
-------
576585
out: array
577-
an array of shape ``(n,)`` containing the sample frequencies. The returned array must have the default real-valued floating-point data type.
586+
an array of shape ``(n,)`` containing the sample frequencies.
578587
579588
Notes
580589
-----
@@ -586,7 +595,14 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> ar
586595
"""
587596

588597

589-
def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> array:
598+
def rfftfreq(
599+
n: int,
600+
/,
601+
*,
602+
d: float = 1.0,
603+
dtype: Optional[dtype] = None,
604+
device: Optional[device] = None,
605+
) -> array:
590606
"""
591607
Computes the discrete Fourier transform sample frequencies (for ``rfft`` and ``irfft``).
592608
@@ -605,13 +621,15 @@ def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[device] = None) -> a
605621
window length.
606622
d: float
607623
sample spacing between individual samples of the Fourier transform input. Default: ``1.0``.
624+
dtype: Optional[dtype]
625+
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``.
608626
device: Optional[device]
609627
device on which to place the created array. Default: ``None``.
610628
611629
Returns
612630
-------
613631
out: array
614-
an array of shape ``(n//2+1,)`` containing the sample frequencies. The returned array must have the default real-valued floating-point data type.
632+
an array of shape ``(n//2+1,)`` containing the sample frequencies.
615633
616634
Notes
617635
-----

0 commit comments

Comments
 (0)