Implement test_fftn and test_ifftn

Author: honno

array_api_tests/shape_helpers.py

@@ -1,6 +1,6 @@
 import math
 from itertools import product
-from typing import Iterator, List, Optional, Tuple, Union
+from typing import Iterator, List, Optional, Sequence, Tuple, Union
 
 from ndindex import iter_indices as _iter_indices
 
@@ -66,10 +66,12 @@ def broadcast_shapes(*shapes: Shape):
 
 
 def normalise_axis(
-    axis: Optional[Union[int, Tuple[int, ...]]], ndim: int
+    axis: Optional[Union[int, Sequence[int]]], ndim: int
 ) -> Tuple[int, ...]:
     if axis is None:
         return tuple(range(ndim))
+    elif isinstance(axis, Sequence) and not isinstance(axis, tuple):
+        axis = tuple(axis)
     axes = axis if isinstance(axis, tuple) else (axis,)
     axes = tuple(axis if axis >= 0 else ndim + axis for axis in axes)
     return axes

array_api_tests/test_fft.py

@@ -1,5 +1,5 @@
 import math
-from typing import Optional
+from typing import List, Optional
 
 import pytest
 from hypothesis import given
@@ -10,6 +10,7 @@
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
+from . import shape_helpers as sh
 from . import xps
 from ._array_module import mod as xp
 
@@ -23,9 +24,9 @@
 fft_shapes_strat = hh.shapes(min_dims=1).filter(lambda s: math.prod(s) > 1)
 
 
-def n_axis_norm_kwargs(x: Array, data: st.DataObject) -> tuple:
+def draw_n_axis_norm_kwargs(x: Array, data: st.DataObject) -> tuple:
     size = math.prod(x.shape)
-    n = data.draw(st.none() | st.integers(size // 2, size * 2), label="n")
+    n = data.draw(st.none() | st.integers((size // 2), math.ceil(size * 1.5)), label="n")
     axis = data.draw(st.integers(-1, x.ndim - 1), label="axis")
     norm = data.draw(st.sampled_from(["backward", "ortho", "forward"]), label="norm")
     kwargs = data.draw(
@@ -39,6 +40,32 @@ def n_axis_norm_kwargs(x: Array, data: st.DataObject) -> tuple:
     return n, axis, norm, kwargs
 
 
+def draw_s_axes_norm_kwargs(x: Array, data: st.DataObject) -> tuple:
+    all_axes = list(range(x.ndim))
+    axes = data.draw(
+        st.none() | st.lists(st.sampled_from(all_axes), min_size=1, unique=True),
+        label="axes",
+    )
+    _axes = all_axes if axes is None else axes
+    axes_sides = [x.shape[axis] for axis in _axes]
+    s_strat = st.tuples(
+        *[st.integers(max(side // 2, 1), math.ceil(side * 1.5)) for side in axes_sides]
+    )
+    if axes is None:
+        s_strat = st.none() | s_strat
+    s = data.draw(s_strat, label="s")
+    norm = data.draw(st.sampled_from(["backward", "ortho", "forward"]), label="norm")
+    kwargs = data.draw(
+        hh.specified_kwargs(
+            ("s", s, None),
+            ("axes", axes, None),
+            ("norm", norm, "backward"),
+        ),
+        label="kwargs",
+    )
+    return s, axes, norm, kwargs
+
+
 def assert_fft_dtype(func_name: str, *, in_dtype: DataType, out_dtype: DataType):
     if in_dtype == xp.float32:
         expected = xp.complex64
@@ -63,12 +90,32 @@ def assert_n_axis_shape(
     ph.assert_shape(func_name, out_shape=out.shape, expected=expected_shape)
 
 
+def assert_s_axes_shape(
+    func_name: str,
+    *,
+    x: Array,
+    s: Optional[List[int]],
+    axes: Optional[List[int]],
+    out: Array,
+):
+    _axes = sh.normalise_axis(axes, x.ndim)
+    _s = x.shape if s is None else s
+    expected = []
+    for i in range(x.ndim):
+        if i in _axes:
+            side = _s[_axes.index(i)]
+        else:
+            side = x.shape[i]
+        expected.append(side)
+    ph.assert_shape(func_name, out_shape=out.shape, expected=tuple(expected))
+
+
 @given(
     x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat),
     data=st.data(),
 )
 def test_fft(x, data):
-    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+    n, axis, norm, kwargs = draw_n_axis_norm_kwargs(x, data)
 
     out = xp.fft.fft(x, **kwargs)
 
@@ -81,25 +128,46 @@ def test_fft(x, data):
     data=st.data(),
 )
 def test_ifft(x, data):
-    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+    n, axis, norm, kwargs = draw_n_axis_norm_kwargs(x, data)
 
     out = xp.fft.ifft(x, **kwargs)
 
     assert_fft_dtype("ifft", in_dtype=x.dtype, out_dtype=out.dtype)
     assert_n_axis_shape("ifft", x=x, n=n, axis=axis, out=out)
 
 
-# TODO:
-# test_fftn
-# test_ifftn
+@given(
+    x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat),
+    data=st.data(),
+)
+def test_fftn(x, data):
+    s, axes, norm, kwargs = draw_s_axes_norm_kwargs(x, data)
+
+    out = xp.fft.fftn(x, **kwargs)
+
+    assert_fft_dtype("fftn", in_dtype=x.dtype, out_dtype=out.dtype)
+    assert_s_axes_shape("fftn", x=x, s=s, axes=axes, out=out)
+
+
+@given(
+    x=xps.arrays(dtype=hh.all_floating_dtypes(), shape=fft_shapes_strat),
+    data=st.data(),
+)
+def test_ifftn(x, data):
+    s, axes, norm, kwargs = draw_s_axes_norm_kwargs(x, data)
+
+    out = xp.fft.ifftn(x, **kwargs)
+
+    assert_fft_dtype("ifftn", in_dtype=x.dtype, out_dtype=out.dtype)
+    assert_s_axes_shape("ifftn", x=x, s=s, axes=axes, out=out)
 
 
 @given(
     x=xps.arrays(dtype=xps.floating_dtypes(), shape=fft_shapes_strat),
     data=st.data(),
 )
 def test_rfft(x, data):
-    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+    n, axis, norm, kwargs = draw_n_axis_norm_kwargs(x, data)
 
     out = xp.fft.rfft(x, **kwargs)
 
@@ -112,7 +180,7 @@ def test_rfft(x, data):
     data=st.data(),
 )
 def test_irfft(x, data):
-    n, axis, norm, kwargs = n_axis_norm_kwargs(x, data)
+    n, axis, norm, kwargs = draw_n_axis_norm_kwargs(x, data)
 
     out = xp.fft.irfft(x, **kwargs)