Add FFT functions

unumpy/fft/__init__.py

@@ -0,0 +1 @@
+from ._multimethods import *

unumpy/fft/_multimethods.py

@@ -0,0 +1,116 @@
+import functools
+import operator
+from uarray import create_multimethod, mark_as, all_of_type, Dispatchable
+import builtins
+
+create_numpy = functools.partial(create_multimethod, domain="numpy")
+
+from .._multimethods import ndarray, _identity_argreplacer, _self_argreplacer
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def fft(a, n=None, axis=-1, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def ifft(a, n=None, axis=-1, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def fft2(a, s=None, axes=(-2, -1), norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def ifft2(a, s=None, axes=(-2, -1), norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def fftn(a, s=None, axes=None, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def ifftn(a, s=None, axes=None, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def rfft(a, n=None, axis=-1, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def irfft(a, n=None, axis=-1, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def rfft2(a, s=None, axes=(-2, -1), norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def irfft2(a, s=None, axes=(-2, -1), norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def rfftn(a, s=None, axes=None, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def irfftn(a, s=None, axes=None, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def hfft(a, n=None, axis=-1, norm=None):
+    return (a,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def ihfft(a, n=None, axis=-1, norm=None):
+    return (a,)
+
+
+@create_numpy(_identity_argreplacer)
+@all_of_type(ndarray)
+def fftfreq(n, d=1.0):
+    return ()
+
+
+@create_numpy(_identity_argreplacer)
+@all_of_type(ndarray)
+def rfftfreq(n, d=1.0):
+    return ()
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def fftshift(x, axes=None):
+    return (x,)
+
+
+@create_numpy(_self_argreplacer)
+@all_of_type(ndarray)
+def ifftshift(x, axes=None):
+    return (x,)

unumpy/numpy_backend.py

@@ -1,6 +1,6 @@
 import numpy as np
 from uarray import Dispatchable, wrap_single_convertor
-from unumpy import ufunc, ufunc_list, ndarray, dtype
+from unumpy import ufunc, ufunc_list, ndarray, dtype, fft
 import unumpy
 import functools
 
@@ -24,11 +24,13 @@ def __ua_function__(method, args, kwargs):
     if method in _implementations:
         return _implementations[method](*args, **kwargs)
 
-    if not hasattr(np, method.__name__):
+    if hasattr(np, method.__name__):
+        return getattr(np, method.__name__)(*args, **kwargs)
+    elif hasattr(np.linalg, method.__name__):
+        return getattr(np.linalg, method.__name__)(*args, **kwargs)
+    else:
         return NotImplemented
 
-    return getattr(np, method.__name__)(*args, **kwargs)
-
 
 @wrap_single_convertor
 def __ua_convert__(value, dispatch_type, coerce):

unumpy/tests/test_numpy.py

@@ -281,3 +281,41 @@ def test_array_creation(backend, method, args, kwargs):
         assert ret.dtype == ndt(dtype)
     else:
         assert ret.dtype == dtype
+
+
+@pytest.mark.parametrize(
+    "method, args, kwargs",
+    [
+        # (np.fft.fft, (4 * np.eye(4),), {}),
+        (np.fft.ifft, ([[1, 2], [3, 4]],), {}),
+        (np.fft.fft2, (4 * np.eye(4),), {}),
+        (np.fft.ifft2, (4 * np.eye(4),), {}),
+        (np.fft.fftn, (4 * np.eye(4),), {}),
+        (np.fft.ifftn, (4 * np.eye(4),), {}),
+        (np.fft.rfft, (4 * np.eye(4),), {}),
+        (np.fft.irfft, (4 * np.eye(4),), {}),
+        (np.fft.rfft2, (4 * np.eye(4),), {}),
+        (np.fft.irfft2, (4 * np.eye(4),), {}),
+        (np.fft.rfftn, (4 * np.eye(4),), {}),
+        (np.fft.irfftn, (4 * np.eye(4)), {}),
+        (np.fft.hfft, (4 * np.eye(4)), {}),
+        (np.fft.ihfft, (4 * np.eye(4)), {}),
+        (np.fft.rfftn, (4 * np.eye(4)), {}),
+        (np.fft.fftfreq, (10), {}),
+        (np.fft.rfftfreq, (10), {}),
+        (np.fft.fftshift, (4 * np.eye(4),), {}),
+        (np.fft.ifftshift, (4 * np.eye(4),), {}),
+    ],
+)
+def test_fft(backend, method, args, kwargs):
+    backend, types = backend
+    try:
+        with ua.set_backend(backend, coerce=True):
+            ret = method(*args, **kwargs)
+    except ua.BackendNotImplementedError:
+        if backend in FULLY_TESTED_BACKENDS and (backend, method) not in EXCEPTIONS:
+            raise
+        pytest.xfail(reason="The backend has no implementation for this ufunc.")
+
+    if isinstance(ret, da.Array):
+        ret.compute()