Add fft support for numpy and cupy

[asmeurer]

This is based off of numpy/numpy#25317

Still need to

• Add support for torch
• Test cupy
• Make sure everything here is correct
• Update tests

[asmeurer]

Need to re-run the tests here, but NumPy 2.0 shouldn't require any wrapping for fft. For NumPy 1.26, the only wrapping that should be needed is to avoid upcasting 32-bit dtypes. CuPy doesn't have that bug, so presumably no wrapping will be needed there (although I haven't run the tests yet so they may have other issues).

[asmeurer]

Confirmed numpy main passes all fft tests in the test suite.

[asmeurer]

CuPy 13.0.0 also passes all fft tests (although annoyingly you have to use array_api_compat.cupy because cupy.bool doesn't exist, CC @leofang).

[leofang]

It's bool_ not bool 😉
https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.bool_
but how is bool relevant to fft?

[asmeurer]

It means you can't run the test suite against cupy:

$ARRAY_API_TESTS_MODULE=cupy pytest array_api_tests
ImportError while loading conftest '/home/aaronmeurer/Documents/array-api-tests/conftest.py'.
conftest.py:13: in <module>
    from reporting import pytest_metadata, pytest_json_modifyreport, add_extra_json_metadata # noqa
reporting.py:1: in <module>
    from array_api_tests.dtype_helpers import dtype_to_name
array_api_tests/dtype_helpers.py:138: in <module>
    all_dtypes = (xp.bool,) + numeric_dtypes
../../miniconda3/envs/array-apis/lib/python3.12/site-packages/cupy/__init__.py:927: in __getattr__
    raise AttributeError(f"module 'cupy' has no attribute {name!r}")
E   AttributeError: module 'cupy' has no attribute 'bool'. Did you mean: 'bool_'?

you have to go through the compat layer instead (ARRAY_API_TESTS_MODULE=array_api_compat.cupy).

(note that the standard uses bool, not bool_, and also this will be changing in numpy 2.0)

[asmeurer]

With that being said, I feel like this is something that we fixed previously. Is this a regression @honno?

[leofang]

Ah OK, but CuPy main namespace is not array-api compliant yet. When CuPy becomes fully NumPy 2.0 compatible we should get the fix for free (cc: @kmaehashi).

[honno]

With that being said, I feel like this is something that we fixed previously. Is this a regression @honno?

I think we did do some work to make testing the top-level numpy namespace work, although as we know that does have a np.bool (which aliases to the builtin bool) so only threw us warnings and didn't really prevent testing 🤔

Maybe we could special case things like this in the test suite, but if there's momentum for CuPy to have a cupy.bool then I'm guessing it isn't quite worth it.

[asmeurer]

The latest version of numpy also has this problem:

ImportError while loading conftest '/Users/aaronmeurer/Documents/array-api-tests/conftest.py'.
conftest.py:13: in <module>
    from reporting import pytest_metadata, pytest_json_modifyreport, add_extra_json_metadata # noqa
reporting.py:1: in <module>
    from array_api_tests.dtype_helpers import dtype_to_name
array_api_tests/dtype_helpers.py:138: in <module>
    all_dtypes = (xp.bool,) + numeric_dtypes
../../miniconda3/envs/array-apis/lib/python3.11/site-packages/numpy/__init__.py:324: in __getattr__
    raise AttributeError(__former_attrs__[attr])
E   AttributeError: module 'numpy' has no attribute 'bool'.
E   `np.bool` was a deprecated alias for the builtin `bool`. To avoid this error in existing code, use `bool` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.bool_` here.
E   The aliases was originally deprecated in NumPy 1.20; for more details and guidance see the original release note at:
E       https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations

[asmeurer]

CuPy test failures:

____________________________________________________________________________________________________ test_fftn ____________________________________________________________________________________________________

    @given(x=hh.arrays(dtype=xps.complex_dtypes(), shape=fft_shapes_strat), data=st.data())
>   def test_fftn(x, data):

array_api_tests/test_fft.py:153: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
array_api_tests/test_fft.py:159: in test_fftn
    assert_s_axes_shape("fftn", x=x, s=s, axes=axes, out=out)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

func_name = 'fftn'

    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))
E       AssertionError: out.shape=(3, 4), but should be (4, 3) [fftn()]
E       Falsifying example: test_fftn(
E           x=array([[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
E                  [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
E                  [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]], dtype=complex64),
E           data=data(...),
E       )
E       Draw 1 (axes): [1, 0]
E       Draw 2 (s): (3, 4)
E       Draw 3 (norm): 'ortho'
E       Draw 4 (kwargs): {'s': (3, 4), 'axes': [1, 0], 'norm': 'ortho'}

array_api_tests/test_fft.py:129: AssertionError
___________________________________________________________________________________________________ test_rfftn ____________________________________________________________________________________________________

    @given(x=hh.arrays(dtype=xps.floating_dtypes(), shape=fft_shapes_strat), data=st.data())
>   def test_rfftn(x, data):

array_api_tests/test_fft.py:212: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

x = array([[0., 0.],
       [0., 0.],
       [0., 0.]], dtype=float32), data = data(...)

    @given(x=hh.arrays(dtype=xps.floating_dtypes(), shape=fft_shapes_strat), data=st.data())
    def test_rfftn(x, data):
        s, axes, norm, kwargs = draw_s_axes_norm_kwargs(x, data)
    
        out = xp.fft.rfftn(x, **kwargs)
    
        assert_float_to_complex_dtype("rfftn", in_dtype=x.dtype, out_dtype=out.dtype)
    
        _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)
        expected[_axes[-1]] = _s[-1] // 2 + 1
>       ph.assert_shape("rfftn", out_shape=out.shape, expected=tuple(expected))
E       AssertionError: out.shape=(2, 2), but should be (2, 3) [rfftn()]
E       Falsifying example: test_rfftn(
E           x=array([[0., 0.],
E                  [0., 0.],
E                  [0., 0.]], dtype=float32),
E           data=data(...),
E       )
E       Draw 1 (axes): [1, 0]
E       Draw 2 (s): (3, 2)
E       Draw 3 (norm): 'ortho'
E       Draw 4 (kwargs): {'s': (3, 2), 'axes': [1, 0], 'norm': 'ortho'}

array_api_tests/test_fft.py:229: AssertionError

[asmeurer]

CC @leofang

>>> import numpy as np
>>> x = np.zeros((3, 2))
>>> np.fft.rfftn(x, s=(3, 2), axes=(1, 0), norm='ortho')
array([[0.+0.j, 0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j, 0.+0.j]])
>>> np.fft.rfftn(x, s=(3, 2), axes=(1, 0), norm='ortho').shape
(2, 3)
>>> import cupy as cp
>>> x = cp.zeros((3, 2))
>>> cp.fft.rfftn(x, s=(3, 2), axes=(1, 0), norm='ortho').shape
(2, 2)
>>> cp.fft.rfftn(x, s=(3, 2), axes=(1, 0), norm='ortho')
array([[0.+0.j, 0.+0.j],
       [0.+0.j, 0.+0.j]])

and a similar issue exists for fftn.

[asmeurer]

Actually the norm argument is irrelevant. Hypothesis seems to have failed at shrinking that out properly.

[asmeurer]

ifftn also appears to be wrong

E       AssertionError: out.shape=(2, 1, 5), but should be (1, 2, 5) [ifftn()]
E       Falsifying example: test_ifftn(
E           x=array([[[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]],
E           
E                  [[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]]], dtype=complex64),
E           data=data(...),
E       )
E       Draw 1 (axes): [1, 0, 2]
E       Draw 2 (s): (2, 1, 5)
E       Draw 3 (norm): 'ortho'
E       Draw 4 (kwargs): {'s': (2, 1, 5), 'axes': [1, 0, 2], 'norm': 'ortho'}

irfftn seems to be fine, though.