Repository: google/jax

Issue:
numpy fft not yet implemented
As a fairly fundamental operation I was surprised to find out this wasn't implemented.

NotImplementedError: Numpy function <function fft at 0x7fdb56275510> not yet implemented

In the short term is there a suggested workaround? I suppose I could use a full dft matrix.


Discussion:
We haven't implemented any FFT functionality yet. It's probably best that we just implement it. We're trying to be fairly user-driven in the things that we add.

Do you just need to compute FFTs (easy enough), or do you also need to differentiate through them?
Being able to differentiate through the fft and ifft would be amazing. In particular for end-to-end optimization of frequency-domain audio processing algorithms.
I would like to second this functionality, it will make a lot of signal processing users happy.
I started working on this earlier and got distracted by other things, I'll try to get something up for review soon!

Resolving patch:
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -19,4 +19,5 @@
                   _reduce_window_min, _reduce_window_prod, _float, _complex,
                   _input_dtype, _const, _eq_meet, _safe_mul, _abstractify)
 from .lax_control_flow import *
+from .lax_fft import *
 from .lax_parallel import *
diff --git a/jax/lax/lax_fft.py b/jax/lax/lax_fft.py
new file mode 100644
--- /dev/null
+++ b/jax/lax/lax_fft.py
@@ -0,0 +1,48 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from jax.abstract_arrays import ShapedArray
+from jax.core import Primitive
+from jax.interpreters import xla
+from ..interpreters import ad
+
+
+def fft(x, fft_type, fft_lengths=None):
+  if fft_lengths is None:
+    fft_lengths = x.shape
+  else:
+    fft_lengths = tuple(fft_lengths)
+  return fft_p.bind(x, fft_type=fft_type, fft_lengths=fft_lengths)
+
+def fft_impl(x, fft_type, fft_lengths):
+  return xla.apply_primitive(fft_p, x, fft_type=fft_type, fft_lengths=fft_lengths)
+
+def fft_abstract_eval(x, fft_type, fft_lengths):
+  return ShapedArray(x.shape, x.dtype)
+
+def fft_translation_rule(c, x, fft_type, fft_lengths):
+  return c.Fft(x, fft_type, fft_lengths)
+
+def fft_transpose_rule(t, fft_type, fft_lengths):
+  return fft(t, fft_type, fft_lengths),
+
+fft_p = Primitive('fft')
+fft_p.def_impl(fft_impl)
+fft_p.def_abstract_eval(fft_abstract_eval)
+xla.translations[fft_p] = fft_translation_rule
+ad.deflinear(fft_p, fft_transpose_rule)
diff --git a/jax/numpy/fft.py b/jax/numpy/fft.py
--- a/jax/numpy/fft.py
+++ b/jax/numpy/fft.py
@@ -18,8 +18,56 @@
 
 import numpy as onp
 
+from .. import lax
+from ..lib.xla_bridge import xla_client, canonicalize_dtype
 from ..util import get_module_functions
 from .lax_numpy import _not_implemented
+from .lax_numpy import _wraps
+from . import lax_numpy as np
+
+
+def _promote_to_complex(arg):
+  dtype = np.result_type(arg, onp.complex64)
+  # XLA's FFT op only supports C64.
+  if dtype == onp.complex128:
+    dtype = onp.complex64
+  return lax.convert_element_type(arg, dtype)
+
+@_wraps(onp.fft.fftn)
+def fftn(a, s=None, axes=None, norm=None):
+  # TODO(skye): implement padding/cropping based on 's'.
+  if s is not None:
+    raise NotImplementedError("jax.np.fftn only supports s=None, got %s" % s)
+  if norm is not None:
+    raise NotImplementedError("jax.np.fftn only supports norm=None, got %s" % norm)
+  if s is not None and axes is not None and len(s) != len(axes):
+    # Same error as numpy.
+    raise ValueError("Shape and axes have different lengths.")
+
+  orig_axes = axes
+  if axes is None:
+    if s is None:
+      axes = range(a.ndim)
+    else:
+      axes = range(a.ndim - len(s), a.ndim)
+
+  # XLA doesn't support 0-rank axes.
+  if len(axes) == 0:
+    return a
+
+  if len(axes) != len(set(axes)):
+    raise ValueError(
+        "jax.np.fftn does not support repeated axes. Got axes %s." % axes)
+
+  if any(axis in range(a.ndim - 3) for axis in axes):
+    raise ValueError(
+        "jax.np.fftn only supports 1D, 2D, and 3D FFTs over the innermost axes."
+        " Got axes %s with input rank %s." % (orig_axes, a.ndim))
+
+  if s is None:
+    s = [a.shape[axis] for axis in axes]
+  a = _promote_to_complex(a)
+  return lax.fft(a, xla_client.FftType.FFT, s)
 
 for func in get_module_functions(onp.fft):
   if func.__name__ not in globals():


Test patch:
diff --git a/tests/fft_test.py b/tests/fft_test.py
new file mode 100644
--- /dev/null
+++ b/tests/fft_test.py
@@ -0,0 +1,90 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as onp
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from jax import numpy as np
+from jax import test_util as jtu
+
+float_dtypes = [onp.float32, onp.float64]
+complex_dtypes = [onp.complex64, onp.complex128]
+inexact_dtypes = float_dtypes + complex_dtypes
+int_dtypes = [onp.int32, onp.int64]
+bool_dtypes = [onp.bool_]
+all_dtypes = float_dtypes + complex_dtypes + int_dtypes + bool_dtypes
+
+
+def _get_fftn_test_axes(shape):
+  axes = [[]]
+  ndims = len(shape)
+  # XLA's FFT op only supports up to 3 innermost dimensions.
+  if ndims <= 3: axes.append(None)
+  for naxes in range(1, min(ndims, 3) + 1):
+    axes.append(range(ndims - naxes, ndims))
+  return axes
+
+
+class FftTest(jtu.JaxTestCase):
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_shape={}_axes={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), axes),
+       "axes": axes, "shape": shape, "dtype": dtype, "rng": rng}
+      for rng in [jtu.rand_default()]
+      for dtype in all_dtypes
+      for shape in [(10,), (10, 10), (2, 3, 4), (2, 3, 4, 5)]
+      for axes in _get_fftn_test_axes(shape)))
+  def testFftn(self, shape, dtype, axes, rng):
+    args_maker = lambda: (rng(shape, dtype),)
+    np_fn = lambda a: np.fft.fftn(a, axes=axes)
+    onp_fn = lambda a: onp.fft.fftn(a, axes=axes)
+    self._CheckAgainstNumpy(onp_fn, np_fn, args_maker, check_dtypes=True,
+                            tol=1e-4)
+    self._CompileAndCheck(np_fn, args_maker, check_dtypes=True)
+    # Test gradient for differentiable types.
+    if dtype in inexact_dtypes:
+      # TODO(skye): can we be more precise?
+      tol = 1e-1
+      jtu.check_grads(np_fn, args_maker(), order=1, atol=tol, rtol=tol)
+      jtu.check_grads(np_fn, args_maker(), order=2, atol=tol, rtol=tol)
+
+  def testFftnErrors(self):
+    rng = jtu.rand_default()
+    self.assertRaisesRegexp(
+        ValueError,
+        "jax.np.fftn only supports 1D, 2D, and 3D FFTs over the innermost axes. "
+        "Got axes None with input rank 4.",
+        lambda: np.fft.fftn(rng([2, 3, 4, 5], dtype=onp.float64), axes=None))
+    self.assertRaisesRegexp(
+        ValueError,
+        "jax.np.fftn only supports 1D, 2D, and 3D FFTs over the innermost axes. "
+        "Got axes \[0\] with input rank 4.",
+        lambda: np.fft.fftn(rng([2, 3, 4, 5], dtype=onp.float64), axes=[0]))
+    self.assertRaisesRegexp(
+        ValueError,
+        "jax.np.fftn does not support repeated axes. Got axes \[1, 1\].",
+        lambda: np.fft.fftn(rng([2, 3], dtype=onp.float64), axes=[1, 1]))
+    self.assertRaises(
+        IndexError, lambda: np.fft.fftn(rng([2, 3], dtype=onp.float64), axes=[2]))
+
+
+if __name__ == "__main__":
+  absltest.main()