torch.fft: Multi-dimensional transforms (pytorch#44550)

Summary:
Pull Request resolved: pytorch#44550

Part of the `torch.fft` work (pytorchgh-42175).
This adds n-dimensional transforms: `fftn`, `ifftn`, `rfftn` and `irfftn`.

This is aiming for correctness first, with the implementation on top of the existing `_fft_with_size` restrictions. I plan to follow up later with a more efficient rewrite that makes `_fft_with_size` work with arbitrary numbers of dimensions.

Test Plan: Imported from OSS

Reviewed By: ngimel

Differential Revision: D23846032

Pulled By: mruberry

fbshipit-source-id: e6950aa8be438ec5cb95fb10bd7b8bc9ffb7d824

Author: peterbell10

aten/src/ATen/WrapDimUtils.h

@@ -30,14 +30,15 @@ static inline int64_t maybe_wrap_dim(int64_t dim, const std::vector<std::vector<
   return maybe_wrap_dim(dim, tensor_sizes[0].size());
 }
 
-// wrap each of dims basing on dim_post_expr
-static inline void maybe_wrap_dims(std::vector<int64_t>& dims, int64_t dim_post_expr) {
+// wrap each dim in the dims array, taking dim_post_expr as the true number of dimensions
+static inline void maybe_wrap_dims_n(int64_t* dims, int64_t ndims, int64_t dim_post_expr) {
   if (dim_post_expr <= 0) {
     dim_post_expr = 1; // this will make range [-1, 0]
   }
   int64_t min = -dim_post_expr;
   int64_t max = dim_post_expr - 1;
-  for (auto& dim : dims) {
+  for (int64_t i = 0; i < ndims; ++i) {
+    auto &dim = dims[i];
     if (dim < min || dim > max) {
       TORCH_CHECK_INDEX(false,
         "Dimension out of range (expected to be in range of [",
@@ -47,6 +48,13 @@ static inline void maybe_wrap_dims(std::vector<int64_t>& dims, int64_t dim_post_
   }
 }
 
+// Wrap each dim in a contiguous container, taking dim_post_expr as the true number of dimensions
+// E.g. could also be std::array or c10::SmallVector
+template <typename Container>
+inline void maybe_wrap_dims(Container& dims, int64_t dim_post_expr) {
+  return maybe_wrap_dims_n(dims.data(), dims.size(), dim_post_expr);
+}
+
 // previously, size [0] tensors were the only possible empty tensors; thus, it wasn't possible
 // to cat empty tensors unless all the other tensors were 1-dimensional, so we allowed these tensors
 // to be "skipped" (both for wrap dimension behavior and dimension size checking).

aten/src/ATen/native/SpectralOps.cpp

@@ -203,6 +203,116 @@ Tensor fft_c2c(Tensor input, c10::optional<int64_t> n_opt,
   return out;
 }
 
+// Dimensions to transform, and the signal shape in those dimensions
+struct ShapeAndDims {
+  DimVector shape, dim;
+};
+
+// Pre-process n-dimensional fft's `s` and `dim` arguments.
+// Wraps dimensions and applies defaulting behavior.
+// Also checks transform dims are unique and transform shape is non-empty.
+ShapeAndDims canonicalize_fft_shape_and_dim_args(
+    Tensor input, c10::optional<IntArrayRef> shape, c10::optional<IntArrayRef> dim) {
+  const int64_t input_dim = input.dim();
+  const IntArrayRef input_sizes = input.sizes();
+  ShapeAndDims ret;
+
+  if (dim) {
+    ret.dim.resize(dim->size());
+    std::copy(dim->begin(), dim->end(), ret.dim.begin());
+    maybe_wrap_dims(ret.dim, input_dim);
+
+    // Check dims are unique
+    DimVector copy = ret.dim;
+    std::sort(copy.begin(), copy.end());
+    auto duplicate = std::adjacent_find(copy.begin(), copy.end());
+    TORCH_CHECK(duplicate == copy.end(), "FFT dims must be unique");
+  }
+
+  if (shape) {
+    // Has shape, may have dim
+    TORCH_CHECK(!dim || dim->size() == shape->size(),
+                "When given, dim and shape arguments must have the same length");
+    TORCH_CHECK(shape->size() <= input_dim,
+                "Got shape with ", shape->size(), " values but input tensor "
+                "only has ", input_dim, " dimensions.");
+    const int64_t transform_ndim = shape->size();
+    // If shape is given, dims defaults to the last shape.size() dimensions
+    if (!dim) {
+      ret.dim.resize(transform_ndim);
+      std::iota(ret.dim.begin(), ret.dim.end(), input_dim - transform_ndim);
+    }
+
+    // Translate shape of -1 to the default length
+    ret.shape.resize(transform_ndim);
+    for (int64_t i = 0; i < transform_ndim; ++i) {
+      const auto n = (*shape)[i];
+      ret.shape[i] = n == -1 ? input_sizes[ret.dim[i]] : n;
+    }
+  } else if (!dim) {
+    // No shape, no dim
+    ret.dim.resize(input_dim);
+    std::iota(ret.dim.begin(), ret.dim.end(), int64_t{0});
+    ret.shape.resize(input_dim);
+    std::copy(input_sizes.begin(), input_sizes.end(), ret.shape.begin());
+  } else {
+    // No shape, has dim
+    ret.shape.resize(ret.dim.size());
+    for (int64_t i = 0; i < ret.dim.size(); ++i) {
+      ret.shape[i] = input_sizes[ret.dim[i]];
+    }
+  }
+  
+  for (int64_t i = 0; i < ret.shape.size(); ++i) {
+    TORCH_CHECK(ret.shape[i] > 0,
+                "Invalid number of data points (", ret.shape[i], ") specified");
+  }
+  
+  return ret;
+}
+
+// Complex to complex n-dimensional fft
+Tensor fftn_c2c(
+    const Tensor& input, IntArrayRef shape, IntArrayRef dim,
+    c10::optional<std::string> norm_str, bool forward) {
+  TORCH_CHECK(input.is_complex(), "Expected a complex input tensor to FFT");
+  const auto input_dim = input.dim();
+
+  Tensor x = resize_fft_input(input, dim, shape);
+  x = at::view_as_real(x);
+
+  const int64_t transform_ndim = dim.size();
+  const auto norm = norm_from_string(norm_str, forward);
+  // _fft_with_size only supports 3 dimensions being transformed at a time.
+  // This limit is inherited from cuFFT.
+  constexpr int64_t max_signal_ndim = 3;
+
+  // Transform n dimensions, up to 3 at a time
+  // TODO: rewrite _fft_with_size to transform more than 3 dimensions at once.
+  for (int64_t i = 0; i < transform_ndim; i += max_signal_ndim) {
+    const int64_t signal_ndim = std::min(transform_ndim - i, max_signal_ndim);
+    DimVector source_dim(signal_ndim);
+    DimVector dest_dim(signal_ndim);
+
+    for (int64_t j = 0; j < signal_ndim; ++j) {
+      source_dim[j] = dim[i + j];
+      dest_dim[j] = j + (input_dim - signal_ndim);
+    }
+
+    // _fft operates on up-to the last 3 dims, so move selected dims to the end
+    x = at::movedim(x, source_dim, dest_dim);
+
+    x = _fft(x, signal_ndim, /*complex_input=*/true, /*complex_output=*/true,
+             /*inverse=*/!forward, /*signal_sizes=*/{}, /*normalization=*/norm,
+             /*onesided=*/false);
+
+    // Move transform dims back to their original order
+    x = at::movedim(x, dest_dim, source_dim);
+  }
+
+  return at::view_as_complex(x);
+}
+
 }
 
 // torch.fft.fft, analogous to NumPy's numpy.fft.fft
@@ -240,6 +350,64 @@ Tensor fft_ihfft(const Tensor& self, c10::optional<int64_t> n, int64_t dim,
   return fft_r2c(self, n, dim, norm, /*forward=*/false, /*onesided=*/true);
 }
 
+Tensor fft_fftn(const Tensor& self, c10::optional<IntArrayRef> s,
+                c10::optional<IntArrayRef> dim,
+                c10::optional<std::string> norm) {
+  auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim);
+  // TODO: For real input, perform rfftn then mirror with conjugate symmetry
+  Tensor input = promote_tensor_fft(self, /*require_complex=*/true);
+  return fftn_c2c(input, desc.shape, desc.dim, norm, /*forward=*/true);
+}
+
+Tensor fft_ifftn(const Tensor& self, c10::optional<IntArrayRef> s,
+                c10::optional<IntArrayRef> dim,
+                c10::optional<std::string> norm) {
+  auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim);
+  Tensor input = promote_tensor_fft(self, /*require_complex=*/true);
+  return fftn_c2c(input, desc.shape, desc.dim, norm, /*forward=*/false);
+}
+
+Tensor fft_rfftn(const Tensor& self, c10::optional<IntArrayRef> s,
+                c10::optional<IntArrayRef> dim,
+                c10::optional<std::string> norm) {
+  auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim);
+  TORCH_CHECK(desc.shape.size() > 0, "rfftn must transform at least one axis");
+
+  const auto last_dim = desc.dim.back();
+  const auto last_shape = desc.shape.back();
+  desc.shape.pop_back();
+  desc.dim.pop_back();
+
+  // rfft on last dim to get hermitian complex shape
+  auto x = native::fft_rfft(self, last_shape, last_dim, norm);
+  // Normal fft on remaining dims
+  return fftn_c2c(x, desc.shape, desc.dim, norm, /*forward=*/true);
+}
+
+Tensor fft_irfftn(const Tensor& self, c10::optional<IntArrayRef> s,
+                c10::optional<IntArrayRef> dim,
+                c10::optional<std::string> norm) {
+  auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim);
+  TORCH_CHECK(desc.shape.size() > 0, "irfftn must transform at least one axis");
+
+  const auto last_dim = desc.dim.back();
+  const auto last_shape = [&]() -> c10::optional<int64_t> {
+    // If shape is defaulted in the last dimension,
+    // pass nullopt to irfft and let it calculate the default size
+    if (!s.has_value() || (s->back() == -1)) {
+      return c10::nullopt;
+    }
+    return desc.shape.back();
+  }();
+  desc.shape.pop_back();
+  desc.dim.pop_back();
+
+  // Normal ifft for all but last dim
+  Tensor x = promote_tensor_fft(self, /*require_complex=*/true);
+   x = fftn_c2c(x, desc.shape, desc.dim, norm, /*forward=*/false);
+  // Then 1d irfft on last dim to get real output
+  return native::fft_irfft(x, last_shape, last_dim, norm);
+}
 
 // This is a pass-through wrapper function that does the size check and
 // inferences. The actual forward implementation function is called

aten/src/ATen/native/native_functions.yaml

@@ -7943,6 +7943,26 @@
   use_c10_dispatcher: full
   variants: function
 
+- func: fft_fftn(Tensor self, int[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor
+  python_module: fft
+  use_c10_dispatcher: full
+  variants: function
+
+- func: fft_ifftn(Tensor self, int[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor
+  python_module: fft
+  use_c10_dispatcher: full
+  variants: function
+
+- func: fft_rfftn(Tensor self, int[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor
+  python_module: fft
+  use_c10_dispatcher: full
+  variants: function
+
+- func: fft_irfftn(Tensor self, int[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor
+  python_module: fft
+  use_c10_dispatcher: full
+  variants: function
+
 - func: fft(Tensor self, int signal_ndim, bool normalized=False) -> Tensor
   use_c10_dispatcher: full
   variants: function, method

docs/source/fft.rst

@@ -19,7 +19,11 @@ Functions
 
 .. autofunction:: fft
 .. autofunction:: ifft
+.. autofunction:: fftn
+.. autofunction:: ifftn
 .. autofunction:: rfft
 .. autofunction:: irfft
+.. autofunction:: rfftn
+.. autofunction:: irfftn
 .. autofunction:: hfft
 .. autofunction:: ihfft

test/test_spectral_ops.py

@@ -3,6 +3,7 @@
 import math
 from contextlib import contextmanager
 from itertools import product
+import itertools
 
 from torch.testing._internal.common_utils import \
     (TestCase, run_tests, TEST_NUMPY, TEST_LIBROSA)
@@ -11,7 +12,7 @@
      skipCPUIfNoMkl, skipCUDAIfRocm, deviceCountAtLeast, onlyCUDA)
 
 from distutils.version import LooseVersion
-from typing import Optional
+from typing import Optional, List
 
 
 if TEST_NUMPY:
@@ -115,6 +116,7 @@ def method_fn(t):
 
     @skipCPUIfNoMkl
     @skipCUDAIfRocm
+    @onlyOnCPUAndCUDA
     @unittest.skipIf(not TEST_NUMPY, 'NumPy not found')
     @precisionOverride({torch.complex64: 1e-4, torch.float: 1e-4})
     @dtypes(torch.float, torch.double, torch.complex64, torch.complex128)
@@ -226,11 +228,13 @@ def test_fft_round_trip(self, device, dtype):
     def test_empty_fft(self, device, dtype):
         t = torch.empty(0, device=device, dtype=dtype)
         match = r"Invalid number of data points \([-\d]*\) specified"
-        fft_functions = [torch.fft.fft, torch.fft.ifft, torch.fft.hfft,
-                         torch.fft.irfft]
+        fft_functions = [torch.fft.fft, torch.fft.fftn,
+                         torch.fft.ifft, torch.fft.ifftn,
+                         torch.fft.irfft, torch.fft.irfftn,
+                         torch.fft.hfft]
         # Real-only functions
         if not dtype.is_complex:
-            fft_functions += [torch.fft.rfft, torch.fft.ihfft]
+            fft_functions += [torch.fft.rfft, torch.fft.rfftn, torch.fft.ihfft]
 
         for fn in fft_functions:
             with self.assertRaisesRegex(RuntimeError, match):
@@ -242,6 +246,9 @@ def test_fft_invalid_dtypes(self, device):
         with self.assertRaisesRegex(RuntimeError, "Expected a real input tensor"):
             torch.fft.rfft(t)
 
+        with self.assertRaisesRegex(RuntimeError, "Expected a real input tensor"):
+            torch.fft.rfftn(t)
+
         with self.assertRaisesRegex(RuntimeError, "Expected a real input tensor"):
             torch.fft.ihfft(t)
 
@@ -292,14 +299,17 @@ def test_fft_half_errors(self, device, dtype):
         # TODO: Remove torch.half error when complex32 is fully implemented
         x = torch.randn(64, device=device).to(dtype)
         fft_functions = (torch.fft.fft, torch.fft.ifft,
+                         torch.fft.fftn, torch.fft.ifftn,
                          torch.fft.rfft, torch.fft.irfft,
+                         torch.fft.rfftn, torch.fft.irfftn,
                          torch.fft.hfft, torch.fft.ihfft)
         for fn in fft_functions:
             with self.assertRaisesRegex(RuntimeError, "Unsupported dtype "):
                 fn(x)
 
     @skipCPUIfNoMkl
     @skipCUDAIfRocm
+    @onlyOnCPUAndCUDA
     @dtypes(torch.double, torch.complex128)  # gradcheck requires double
     def test_fft_backward(self, device, dtype):
         test_args = list(product(
@@ -340,6 +350,166 @@ def test_fn(x):
 
                 self.assertTrue(torch.autograd.gradcheck(test_fn, (input,)))
 
+    # nd-fft tests
+
+    @skipCPUIfNoMkl
+    @skipCUDAIfRocm
+    @onlyOnCPUAndCUDA
+    @unittest.skipIf(not TEST_NUMPY, 'NumPy not found')
+    @precisionOverride({torch.complex64: 1e-4, torch.float: 1e-4})
+    @dtypes(torch.float, torch.double, torch.complex64, torch.complex128)
+    def test_fftn_numpy(self, device, dtype):
+        norm_modes = ((None, "forward", "backward", "ortho")
+                      if LooseVersion(np.__version__) >= '1.20.0'
+                      else (None, "ortho"))
+
+        # input_ndim, s, dim
+        transform_desc = [
+            *product(range(2, 5), (None,), (None, (0,), (0, -1))),
+            *product(range(2, 5), (None, (4, 10)), (None,)),
+            (6, None, None),
+            (5, None, (1, 3, 4)),
+            (3, None, (0, -1)),
+            (3, None, (1,)),
+            (1, None, (0,)),
+            (4, (10, 10), None),
+            (4, (10, 10), (0, 1))
+        ]
+
+        fft_functions = ['fftn', 'ifftn', 'irfftn']
+        # Real-only functions
+        if not dtype.is_complex:
+            fft_functions += ['rfftn']
+
+        for input_ndim, s, dim in transform_desc:
+            shape = itertools.islice(itertools.cycle(range(4, 9)), input_ndim)
+            input = torch.randn(*shape, device=device, dtype=dtype)
+            for fname, norm in product(fft_functions, norm_modes):
+                torch_fn = getattr(torch.fft, fname)
+                numpy_fn = getattr(np.fft, fname)
+
+                def fn(t: torch.Tensor, s: Optional[List[int]], dim: Optional[List[int]], norm: Optional[str]):
+                    return torch_fn(t, s, dim, norm)
+
+                torch_fns = (torch_fn, torch.jit.script(fn))
+
+                expected = numpy_fn(input.cpu().numpy(), s, dim, norm)
+                exact_dtype = dtype in (torch.double, torch.complex128)
+                for fn in torch_fns:
+                    actual = fn(input, s, dim, norm)
+                    self.assertEqual(actual, expected, exact_dtype=exact_dtype)
+
+    @skipCUDAIfRocm
+    @skipCPUIfNoMkl
+    @onlyOnCPUAndCUDA
+    @dtypes(torch.float, torch.double, torch.complex64, torch.complex128)
+    def test_fftn_round_trip(self, device, dtype):
+        norm_modes = (None, "forward", "backward", "ortho")
+
+        # input_ndim, dim
+        transform_desc = [
+            *product(range(2, 5), (None, (0,), (0, -1))),
+            *product(range(2, 5), (None,)),
+            (7, None),
+            (5, (1, 3, 4)),
+            (3, (0, -1)),
+            (3, (1,)),
+            (1, 0),
+        ]
+
+        fft_functions = [(torch.fft.fftn, torch.fft.ifftn)]
+
+        # Real-only functions
+        if not dtype.is_complex:
+            fft_functions += [(torch.fft.rfftn, torch.fft.irfftn)]
+
+        for input_ndim, dim in transform_desc:
+            shape = itertools.islice(itertools.cycle(range(4, 9)), input_ndim)
+            x = torch.randn(*shape, device=device, dtype=dtype)
+
+            for (forward, backward), norm in product(fft_functions, norm_modes):
+                if isinstance(dim, tuple):
+                    s = [x.size(d) for d in dim]
+                else:
+                    s = x.size() if dim is None else x.size(dim)
+
+                kwargs = {'s': s, 'dim': dim, 'norm': norm}
+                y = backward(forward(x, **kwargs), **kwargs)
+                # For real input, ifftn(fftn(x)) will convert to complex
+                self.assertEqual(x, y, exact_dtype=(
+                    forward != torch.fft.fftn or x.is_complex()))
+
+    @skipCPUIfNoMkl
+    @skipCUDAIfRocm
+    @onlyOnCPUAndCUDA
+    @dtypes(torch.double, torch.complex128)  # gradcheck requires double
+    def test_fftn_backward(self, device, dtype):
+        # input_ndim, s, dim
+        transform_desc = [
+            *product((2, 3), (None,), (None, (0,), (0, -1))),
+            *product((2, 3), (None, (4, 10)), (None,)),
+            (4, None, None),
+            (3, (10, 10), (0, 1)),
+            (2, (1, 1), (0, 1)),
+            (2, None, (1,)),
+            (1, None, (0,)),
+            (1, (11,), (0,)),
+        ]
+        norm_modes = (None, "forward", "backward", "ortho")
+
+        fft_functions = ['fftn', 'ifftn', 'irfftn']
+        # Real-only functions
+        if not dtype.is_complex:
+            fft_functions += ['rfftn']
+
+        for input_ndim, s, dim in transform_desc:
+            shape = itertools.islice(itertools.cycle(range(4, 9)), input_ndim)
+            input = torch.randn(*shape, device=device, dtype=dtype)
+
+            for fname, norm in product(fft_functions, norm_modes):
+                torch_fn = getattr(torch.fft, fname)
+
+                # Workaround for gradcheck's poor support for complex input
+                # Use real input instead and put view_as_complex into the graph
+                if dtype.is_complex:
+                    def test_fn(x):
+                        return torch_fn(torch.view_as_complex(x), s, dim, norm)
+                    inputs = (torch.view_as_real(input).detach().requires_grad_(),)
+                else:
+                    def test_fn(x):
+                        return torch_fn(x, s, dim, norm)
+                    inputs = (input.detach().requires_grad_(),)
+
+                self.assertTrue(torch.autograd.gradcheck(test_fn, inputs))
+
+    @skipCUDAIfRocm
+    @skipCPUIfNoMkl
+    @onlyOnCPUAndCUDA
+    def test_fftn_invalid(self, device):
+        a = torch.rand(10, 10, 10, device=device)
+        fft_funcs = (torch.fft.fftn, torch.fft.ifftn,
+                     torch.fft.rfftn, torch.fft.irfftn)
+
+        for func in fft_funcs:
+            with self.assertRaisesRegex(RuntimeError, "FFT dims must be unique"):
+                func(a, dim=(0, 1, 0))
+
+            with self.assertRaisesRegex(RuntimeError, "FFT dims must be unique"):
+                func(a, dim=(2, -1))
+
+            with self.assertRaisesRegex(RuntimeError, "dim and shape .* same length"):
+                func(a, s=(1,), dim=(0, 1))
+
+            with self.assertRaisesRegex(IndexError, "Dimension out of range"):
+                func(a, dim=(3,))
+
+            with self.assertRaisesRegex(RuntimeError, "tensor only has 3 dimensions"):
+                func(a, s=(10, 10, 10, 10))
+
+        c = torch.complex(a, a)
+        with self.assertRaisesRegex(RuntimeError, "Expected a real input"):
+            torch.fft.rfftn(c)
+
     # Legacy fft tests
     def _test_fft_ifft_rfft_irfft(self, device, dtype):
         def _test_complex(sizes, signal_ndim, prepro_fn=lambda x: x):

torch/csrc/api/include/torch/fft.h

@@ -35,6 +35,36 @@ inline Tensor ifft(const Tensor& self,
   return torch::fft_ifft(self, n, dim, norm);
 }
 
+/// Computes the N dimensional fast Fourier transform over given dimensions.
+/// See https://pytorch.org/docs/master/fft.html#torch.fft.fftn.
+///
+/// Example:
+/// ```
+/// auto t = torch::randn({128, 128}, dtype=kComplexDouble);
+/// torch::fft::fftn(t);
+/// ```
+inline Tensor fftn(const Tensor& self,
+                   c10::optional<IntArrayRef> s=c10::nullopt,
+                   c10::optional<IntArrayRef> dim=c10::nullopt,
+                   c10::optional<std::string> norm=c10::nullopt) {
+  return torch::fft_fftn(self, s, dim, norm);
+}
+
+/// Computes the N dimensional fast Fourier transform over given dimensions.
+/// See https://pytorch.org/docs/master/fft.html#torch.fft.ifftn.
+///
+/// Example:
+/// ```
+/// auto t = torch::randn({128, 128}, dtype=kComplexDouble);
+/// torch::fft::ifftn(t);
+/// ```
+inline Tensor ifftn(const Tensor& self,
+                   c10::optional<IntArrayRef> s=c10::nullopt,
+                   c10::optional<IntArrayRef> dim=c10::nullopt,
+                   c10::optional<std::string> norm=c10::nullopt) {
+  return torch::fft_ifftn(self, s, dim, norm);
+}
+
 /// Computes the 1 dimensional FFT of real input with onesided Hermitian output.
 /// See https://pytorch.org/docs/master/fft.html#torch.fft.rfft.
 ///
@@ -69,6 +99,36 @@ inline Tensor irfft(const Tensor& self,
   return torch::fft_irfft(self, n, dim, norm);
 }
 
+/// Computes the N dimensional FFT of real input with onesided Hermitian output.
+/// See https://pytorch.org/docs/master/fft.html#torch.fft.rfftn
+///
+/// Example:
+/// ```
+/// auto t = torch::randn({128, 128}, dtype=kDouble);
+/// torch::fft::rfftn(t);
+/// ```
+inline Tensor rfftn(const Tensor& self,
+                    c10::optional<IntArrayRef> s=c10::nullopt,
+                    c10::optional<IntArrayRef> dim=c10::nullopt,
+                    c10::optional<std::string> norm=c10::nullopt) {
+  return torch::fft_rfftn(self, s, dim, norm);
+}
+
+/// Computes the inverse of torch.fft.rfftn.
+/// See https://pytorch.org/docs/master/fft.html#torch.fft.irfftn.
+///
+/// Example:
+/// ```
+/// auto t = torch::randn({128, 128}, dtype=kComplexDouble);
+/// torch::fft::irfftn(t);
+/// ```
+inline Tensor irfftn(const Tensor& self,
+                   c10::optional<IntArrayRef> s=c10::nullopt,
+                   c10::optional<IntArrayRef> dim=c10::nullopt,
+                   c10::optional<std::string> norm=c10::nullopt) {
+  return torch::fft_irfftn(self, s, dim, norm);
+}
+
 /// Computes the 1 dimensional FFT of a onesided Hermitian signal
 ///
 /// The input represents a Hermitian symmetric time domain signal. The returned

torch/fft/__init__.py

@@ -87,6 +87,101 @@
     tensor([0.+0.j, 1.+0.j, 2.+0.j, 3.+0.j])
 """)
 
+fftn = _add_docstr(_fft.fft_fftn, r"""
+fftn(input, s=None, dim=None, norm=None) -> Tensor
+
+Computes the N dimensional discrete Fourier transform of :attr:`input`.
+
+Note:
+
+    The Fourier domain representation of any real signal satisfies the
+    Hermitian property: ``X[i_1, ..., i_n] = conj(X[-i_1, ..., -i_n])``. This
+    function always returns all positive and negative frequency terms even
+    though, for real inputs, half of these values are redundant.
+    :func:`~torch.fft.rfftn` returns the more compact one-sided representation
+    where only the positive frequencies of the last dimension are returned.
+
+Args:
+    input (Tensor): the input tensor
+    s (Tuple[int], optional): Signal size in the transformed dimensions.
+        If given, each dimension ``dim[i]`` will either be zero-padded or
+        trimmed to the length ``s[i]`` before computing the FFT.
+        If a length ``-1`` is specified, no padding is done in that dimension.
+        Default: ``s = [input.size(d) for d in dim]``
+    dim (Tuple[int], optional): Dimensions to be transformed.
+        Default: all dimensions, or the last ``len(s)`` dimensions if :attr:`s` is given.
+    norm (str, optional): Normalization mode. For the forward transform
+        (:func:`~torch.fft.fftn`), these correspond to:
+
+        * ``"forward"`` - normalize by ``1/n``
+        * ``"backward"`` - no normalization
+        * ``"ortho"`` - normalize by ``1/sqrt(n)`` (making the FFT orthonormal)
+
+        Where ``n = prod(s)`` is the logical FFT size.
+        Calling the backward transform (:func:`~torch.fft.ifftn`) with the same
+        normalization mode will apply an overall normalization of ``1/n``
+        between the two transforms. This is required to make
+        :func:`~torch.fft.ifftn` the exact inverse.
+
+        Default is ``"backward"`` (no normalization).
+
+Example:
+
+    >>> import torch.fft
+    >>> x = torch.rand(10, 10, dtype=torch.complex64)
+    >>> fftn = torch.fft.fftn(t)
+
+    The discrete Fourier transform is separable, so :func:`~torch.fft.fftn`
+    here is equivalent to two one-dimensional :func:`~torch.fft.fft` calls:
+
+    >>> two_ffts = torch.fft.fft(torch.fft.fft(x, dim=0), dim=1)
+    >>> torch.allclose(fftn, two_ffts)
+
+""")
+
+ifftn = _add_docstr(_fft.fft_ifftn, r"""
+ifftn(input, s=None, dim=None, norm=None) -> Tensor
+
+Computes the N dimensional inverse discrete Fourier transform of :attr:`input`.
+
+Args:
+    input (Tensor): the input tensor
+    s (Tuple[int], optional): Signal size in the transformed dimensions.
+        If given, each dimension ``dim[i]`` will either be zero-padded or
+        trimmed to the length ``s[i]`` before computing the IFFT.
+        If a length ``-1`` is specified, no padding is done in that dimension.
+        Default: ``s = [input.size(d) for d in dim]``
+    dim (Tuple[int], optional): Dimensions to be transformed.
+        Default: all dimensions, or the last ``len(s)`` dimensions if :attr:`s` is given.
+    norm (str, optional): Normalization mode. For the backward transform
+        (:func:`~torch.fft.ifftn`), these correspond to:
+
+        * ``"forward"`` - no normalization
+        * ``"backward"`` - normalize by ``1/n``
+        * ``"ortho"`` - normalize by ``1/sqrt(n)`` (making the IFFT orthonormal)
+
+        Where ``n = prod(s)`` is the logical IFFT size.
+        Calling the forward transform (:func:`~torch.fft.fftn`) with the same
+        normalization mode will apply an overall normalization of ``1/n`` between
+        the two transforms. This is required to make :func:`~torch.fft.ifftn`
+        the exact inverse.
+
+        Default is ``"backward"`` (normalize by ``1/n``).
+
+Example:
+
+    >>> import torch.fft
+    >>> x = torch.rand(10, 10, dtype=torch.complex64)
+    >>> ifftn = torch.fft.ifftn(t)
+
+    The discrete Fourier transform is separable, so :func:`~torch.fft.ifftn`
+    here is equivalent to two one-dimensional :func:`~torch.fft.ifft` calls:
+
+    >>> two_iffts = torch.fft.ifft(torch.fft.ifft(x, dim=0), dim=1)
+    >>> torch.allclose(ifftn, two_iffts)
+
+""")
+
 rfft = _add_docstr(_fft.fft_rfft, r"""
 rfft(input, n=None, dim=-1, norm=None) -> Tensor
 
@@ -199,6 +294,136 @@
     tensor([0.0000, 1.0000, 2.0000, 3.0000, 4.0000])
 """)
 
+rfftn = _add_docstr(_fft.fft_rfftn, r"""
+rfftn(input, s=None, dim=None, norm=None) -> Tensor
+
+Computes the N-dimensional discrete Fourier transform of real :attr:`input`.
+
+The FFT of a real signal is Hermitian-symmetric,
+``X[i_1, ..., i_n] = conj(X[-i_1, ..., -i_n])`` so the full
+:func:`~torch.fft.fftn` output contains redundant information.
+:func:`~torch.fft.rfftn` instead omits the negative frequencies in the
+last dimension.
+
+Args:
+    input (Tensor): the input tensor
+    s (Tuple[int], optional): Signal size in the transformed dimensions.
+        If given, each dimension ``dim[i]`` will either be zero-padded or
+        trimmed to the length ``s[i]`` before computing the real FFT.
+        If a length ``-1`` is specified, no padding is done in that dimension.
+        Default: ``s = [input.size(d) for d in dim]``
+    dim (Tuple[int], optional): Dimensions to be transformed.
+        Default: all dimensions, or the last ``len(s)`` dimensions if :attr:`s` is given.
+    norm (str, optional): Normalization mode. For the forward transform
+        (:func:`~torch.fft.rfftn`), these correspond to:
+
+        * ``"forward"`` - normalize by ``1/n``
+        * ``"backward"`` - no normalization
+        * ``"ortho"`` - normalize by ``1/sqrt(n)`` (making the real FFT orthonormal)
+
+        Where ``n = prod(s)`` is the logical FFT size.
+        Calling the backward transform (:func:`~torch.fft.irfftn`) with the same
+        normalization mode will apply an overall normalization of ``1/n`` between
+        the two transforms. This is required to make :func:`~torch.fft.irfftn`
+        the exact inverse.
+
+        Default is ``"backward"`` (no normalization).
+
+Example:
+
+    >>> import torch.fft
+    >>> t = torch.rand(10, 10)
+    >>> rfftn = torch.fft.rfftn(t)
+    >>> rfftn.size()
+    torch.Size([10, 6])
+
+    Compared against the full output from :func:`~torch.fft.fftn`, we have all
+    elements up to the Nyquist frequency.
+
+    >>> fftn = torch.fft.fftn(t)
+    >>> torch.allclose(fftn[..., :6], rfftn)
+    True
+
+    The discrete Fourier transform is separable, so :func:`~torch.fft.rfftn`
+    here is equivalent to a combination of :func:`~torch.fft.fft` and
+    :func:`~torch.fft.rfft`:
+
+    >>> two_ffts = torch.fft.fft(torch.fft.rfft(x, dim=1), dim=0)
+    >>> torch.allclose(rfftn, two_ffts)
+
+""")
+
+irfftn = _add_docstr(_fft.fft_irfftn, r"""
+irfftn(input, s=None, dim=None, norm=None) -> Tensor
+
+Computes the inverse of :func:`~torch.fft.rfftn`.
+
+:attr:`input` is interpreted as a one-sided Hermitian signal in the Fourier
+domain, as produced by :func:`~torch.fft.rfftn`. By the Hermitian property, the
+output will be real-valued.
+
+Note:
+    Some input frequencies must be real-valued to satisfy the Hermitian
+    property. In these cases the imaginary component will be ignored.
+    For example, any imaginary component in the zero-frequency term cannot
+    be represented in a real output and so will always be ignored.
+
+Note:
+    The correct interpretation of the Hermitian input depends on the length of
+    the original data, as given by :attr:`s`. This is because each input shape
+    could correspond to either an odd or even length signal. By default, the
+    signal is assumed to be even length and odd signals will not round-trip
+    properly. So, it is recommended to always pass the signal shape :attr:`s`.
+
+Args:
+    input (Tensor): the input tensor
+    s (Tuple[int], optional): Signal size in the transformed dimensions.
+        If given, each dimension ``dim[i]`` will either be zero-padded or
+        trimmed to the length ``s[i]`` before computing the real FFT.
+        If a length ``-1`` is specified, no padding is done in that dimension.
+        Defaults to even output in the last dimension:
+        ``s[-1] = 2*(input.size(dim[-1]) - 1)``.
+    dim (Tuple[int], optional): Dimensions to be transformed.
+        The last dimension must be the half-Hermitian compressed dimension.
+        Default: all dimensions, or the last ``len(s)`` dimensions if :attr:`s` is given.
+    norm (str, optional): Normalization mode. For the backward transform
+        (:func:`~torch.fft.irfftn`), these correspond to:
+
+        * ``"forward"`` - no normalization
+        * ``"backward"`` - normalize by ``1/n``
+        * ``"ortho"`` - normalize by ``1/sqrt(n)`` (making the real IFFT orthonormal)
+
+        Where ``n = prod(s)`` is the logical IFFT size.
+        Calling the forward transform (:func:`~torch.fft.rfftn`) with the same
+        normalization mode will apply an overall normalization of ``1/n`` between
+        the two transforms. This is required to make :func:`~torch.fft.irfftn`
+        the exact inverse.
+
+        Default is ``"backward"`` (normalize by ``1/n``).
+
+Example:
+
+    >>> import torch.fft
+    >>> t = torch.rand(10, 9)
+    >>> T = torch.fft.rfftn(t)
+
+    Without specifying the output length to :func:`~torch.fft.irfft`, the output
+    will not round-trip properly because the input is odd-length in the last
+    dimension:
+
+    >>> torch.fft.irfftn(T).size()
+    torch.Size([10, 10])
+
+    So, it is recommended to always pass the signal shape :attr:`s`.
+
+    >>> roundtrip = torch.fft.irfftn(T, t.size())
+    >>> roundtrip.size()
+    torch.Size([10, 9])
+    >>> torch.allclose(roundtrip, t)
+    True
+
+""")
+
 hfft = _add_docstr(_fft.fft_hfft, r"""
 hfft(input, n=None, dim=-1, norm=None) -> Tensor