Repository: google/jax

Issue:
Feature request: IFFT
Hi there! I'm finding Jax to be very useful in my work, but a new project (spectral Gaussian process stuff) requires me to differentiate through both FFT and IFFT operations. It is possible to express IFFT in terms of FFT, but it would be cleaner to just have an IFFT primitive, I think.

Thank you!

edit: also, vmap support would be very helpful as well. Also for people from the future, here is an implementation using FFT only:

```python
from typing import Optional, Sequence

def ifftn(
    a: np.ndarray,
    s: Sequence[int] = None,
    axes: Sequence[int] = None,
    norm: Optional[str] = None,
) -> np.ndarray:
    """Compute the n-dimensional inverse Fourier transform.

    Args:
        a: Input array.
        s: Shape of the output.
        axes: Axes over which to compute the IFFT.
        norm: Normalization mode, one of ``None`` or ``"ortho"``.

    Returns:
        The n-dimensional IFFT of ``a``.

    """
    if s is not None:
        raise NotImplementedError
    if norm is not None:
        raise NotImplementedError
    n = np.prod(np.array(a.shape)[axes])
    return (1 / n) * np.conj(np.fft.fftn(np.conj(a), axes=axes))

```


Discussion:
@skye fyi

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -354,7 +354,7 @@ def convert_element_type(operand, new_dtype):
     if (onp.issubdtype(old_dtype, onp.complexfloating) and
         not onp.issubdtype(new_dtype, onp.complexfloating)):
       msg = "Casting complex values to real discards the imaginary part"
-      warnings.warn(msg, onp.ComplexWarning)
+      warnings.warn(msg, onp.ComplexWarning, stacklevel=2)
       operand = real(operand)
       old_dtype = _dtype(operand)
     return convert_element_type_p.bind(
diff --git a/jax/numpy/fft.py b/jax/numpy/fft.py
--- a/jax/numpy/fft.py
+++ b/jax/numpy/fft.py
@@ -33,13 +33,14 @@ def _promote_to_complex(arg):
     dtype = onp.complex64
   return lax.convert_element_type(arg, dtype)
 
-@_wraps(onp.fft.fftn)
-def fftn(a, s=None, axes=None, norm=None):
+
+def _fft_core(func_name, fft_type, a, s, axes, norm):
   # TODO(skye): implement padding/cropping based on 's'.
+  full_name = "jax.np.fft." + func_name
   if s is not None:
-    raise NotImplementedError("jax.np.fftn only supports s=None, got %s" % s)
+    raise NotImplementedError("%s only supports s=None, got %s" % (full_name, s))
   if norm is not None:
-    raise NotImplementedError("jax.np.fftn only supports norm=None, got %s" % norm)
+    raise NotImplementedError("%s only supports norm=None, got %s" % (full_name, 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.")
@@ -57,17 +58,28 @@ def fftn(a, s=None, axes=None, norm=None):
 
   if len(axes) != len(set(axes)):
     raise ValueError(
-        "jax.np.fftn does not support repeated axes. Got axes %s." % axes)
+        "%s does not support repeated axes. Got axes %s." % (full_name, 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))
+        "%s only supports 1D, 2D, and 3D FFTs over the innermost axes."
+        " Got axes %s with input rank %s." % (full_name, 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)
+  return lax.fft(a, fft_type, s)
+
+
+@_wraps(onp.fft.fftn)
+def fftn(a, s=None, axes=None, norm=None):
+  return _fft_core('fftn', xla_client.FftType.FFT, a, s, axes, norm)
+
+
+@_wraps(onp.fft.ifftn)
+def ifftn(a, s=None, axes=None, norm=None):
+  return _fft_core('ifftn', xla_client.FftType.IFFT, a, s, axes, norm)
+
 
 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
--- a/tests/fft_test.py
+++ b/tests/fft_test.py
@@ -51,17 +51,21 @@ def _get_fftn_test_axes(shape):
 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}
+      {"testcase_name": "_inverse={}_shape={}_axes={}".format(
+          inverse, jtu.format_shape_dtype_string(shape, dtype), axes),
+       "axes": axes, "shape": shape, "dtype": dtype, "rng": rng,
+       "inverse": inverse}
+      for inverse in [False, True]
       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):
+  def testFftn(self, inverse, 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)
+    np_op = np.fft.ifftn if inverse else np.fft.fftn
+    onp_op = onp.fft.ifftn if inverse else onp.fft.fftn
+    np_fn = lambda a: np_op(a, axes=axes)
+    onp_fn = lambda a: onp_op(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)
@@ -72,24 +76,30 @@ def testFftn(self, shape, dtype, axes, rng):
       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):
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_inverse={}".format(inverse),
+       "inverse": inverse}
+      for inverse in [False, True]))
+  def testFftnErrors(self, inverse):
     rng = jtu.rand_default()
-    self.assertRaisesRegexp(
+    name = 'ifftn' if inverse else 'fftn'
+    func = np.fft.ifftn if inverse else np.fft.fftn
+    self.assertRaisesRegex(
         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(
+        "jax.np.fft.{} only supports 1D, 2D, and 3D FFTs over the innermost axes. "
+        "Got axes None with input rank 4.".format(name),
+        lambda: func(rng([2, 3, 4, 5], dtype=onp.float64), axes=None))
+    self.assertRaisesRegex(
         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(
+        "jax.np.fft.{} only supports 1D, 2D, and 3D FFTs over the innermost axes. "
+        "Got axes \\[0\\] with input rank 4.".format(name),
+        lambda: func(rng([2, 3, 4, 5], dtype=onp.float64), axes=[0]))
+    self.assertRaisesRegex(
         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]))
+        "jax.np.fft.{} does not support repeated axes. Got axes \\[1, 1\\].".format(name),
+        lambda: func(rng([2, 3], dtype=onp.float64), axes=[1, 1]))
     self.assertRaises(
-        IndexError, lambda: np.fft.fftn(rng([2, 3], dtype=onp.float64), axes=[2]))
+        IndexError, lambda: func(rng([2, 3], dtype=onp.float64), axes=[2]))
 
 
 if __name__ == "__main__":