ENH: fft: fftfreq respect default dtype/device pairs

Author: ev-br

array_api_strict/_fft.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from ._array_object import Array
-from ._devices import ALL_DEVICES, Device
+from ._devices import ALL_DEVICES, Device, device_supports_dtype
 from ._data_type_functions import astype
 from ._dtypes import (
     DType,
@@ -14,6 +14,7 @@
     complex64,
     float32,
 )
+from ._info import __array_namespace_info__
 from ._flags import requires_extension
 
 
@@ -269,6 +270,15 @@ def fftfreq(
     np_result = np.fft.fftfreq(n, d=d)
     if dtype:
         np_result = np_result.astype(dtype._np_dtype)
+
+    if not device_supports_dtype(device, DType(np_result.dtype)):
+        if dtype:
+            # user input unsupported
+            raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.")
+
+        dt = __array_namespace_info__().default_dtypes(device=device)["real floating"]
+        np_result = np_result.astype(dt._np_dtype)
+
     return Array._new(np_result, device=device)
 
 @requires_extension('fft')
@@ -293,6 +303,15 @@ def rfftfreq(
     np_result = np.fft.rfftfreq(n, d=d)
     if dtype:
         np_result = np_result.astype(dtype._np_dtype)
+
+    if not device_supports_dtype(device, DType(np_result.dtype)):
+        if dtype:
+            # user input unsupported
+            raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.")
+
+        dt = __array_namespace_info__().default_dtypes(device=device)["real floating"]
+        np_result = np_result.astype(dt._np_dtype)
+
     return Array._new(np_result, device=device)
 
 @requires_extension('fft')

array_api_strict/tests/test_device_support.py

@@ -38,6 +38,18 @@ def test_fft_device_support_real(func_name):
     assert x.device == y.device
 
 
+@pytest.mark.parametrize("func_name", ("fftfreq", "rfftfreq"))
+def test_fft_default_dtype(func_name):
+    func = getattr(xp.fft, func_name)
+    device = xp.Device("F32_device")
+    res = func(3, device=device)
+    assert res.device == device
+    assert res.dtype == xp.__array_namespace_info__().default_dtypes(device=device)["real floating"]
+
+    with pytest.raises(ValueError):
+        func(3, device=device, dtype=xp.float64)
+
+
 class TestF32Device:
     @pytest.mark.parametrize("dtype_str", ["float64", "complex128"])
     def test_f64_raises(self, dtype_str):