Rudimentary test_fft

Author: honno

array_api_tests/_array_module.py

@@ -63,7 +63,7 @@ def __repr__(self):
 _constants = ["e", "inf", "nan", "pi"]
 _funcs = [f.__name__ for funcs in stubs.category_to_funcs.values() for f in funcs]
 _funcs += ["take", "isdtype", "conj", "imag", "real"]  # TODO: bump spec and update array-api-tests to new spec layout
-_top_level_attrs = _dtypes + _constants + _funcs + stubs.EXTENSIONS
+_top_level_attrs = _dtypes + _constants + _funcs + stubs.EXTENSIONS + ["fft"]
 
 for attr in _top_level_attrs:
     try:

array_api_tests/stubs.py

@@ -52,7 +52,7 @@
     all_funcs.extend(funcs)
 name_to_func: Dict[str, FunctionType] = {f.__name__: f for f in all_funcs}
 
-EXTENSIONS: str = ["linalg"]
+EXTENSIONS: List[str] = ["linalg"]  # TODO: add "fft" once stubs available
 extension_to_funcs: Dict[str, List[FunctionType]] = {}
 for ext in EXTENSIONS:
     mod = name_to_mod[ext]

array_api_tests/test_fft.py

@@ -0,0 +1,38 @@
+import math
+
+import pytest
+from hypothesis import given
+
+from array_api_tests.typing import DataType
+
+from . import _array_module as xp
+from . import hypothesis_helpers as hh
+from . import pytest_helpers as ph
+from . import xps
+
+pytestmark = [
+    pytest.mark.ci,
+    pytest.mark.xp_extension("fft"),
+    pytest.mark.min_version("draft"),
+]
+
+
+fft_shapes_strat = hh.shapes(min_dims=1).filter(lambda s: math.prod(s) > 1)
+
+
+def assert_fft_dtype(func_name: str, *, in_dtype: DataType, out_dtype: DataType):
+    if in_dtype == xp.float32:
+        expected = xp.complex64
+    else:
+        assert in_dtype == xp.float64  # sanity check
+        expected = xp.complex128
+    ph.assert_dtype(
+        func_name, in_dtype=in_dtype, out_dtype=out_dtype, expected=expected
+    )
+
+
+@given(x=xps.arrays(dtype=xps.floating_dtypes(), shape=fft_shapes_strat))
+def test_fft(x):
+    out = xp.fft.fft(x)
+    assert_fft_dtype("fft", in_dtype=x.dtype, out_dtype=out.dtype)
+    ph.assert_shape("fft", out_shape=out.shape, expected=x.shape)