C++ API and docs for hfftn (#66127)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66127

cc mruberry peterbell10

Test Plan: Imported from OSS

Reviewed By: dagitses

Differential Revision: D31450216

Pulled By: mruberry

fbshipit-source-id: 2878aee294aa7d74482b66d536258bac0541408d

docs/source/fft.rst

@@ -32,6 +32,10 @@ Fast Fourier Transforms
     irfftn
     hfft
     ihfft
+    hfft2
+    ihfft2
+    hfftn
+    ihfftn
 
 Helper Functions
 ----------------

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

@@ -216,7 +216,7 @@ inline Tensor hfft(const Tensor& self,
 /// Example:
 /// ```
 /// auto T = torch::randn(128, torch::kDouble);
-/// auto t = torch::fft::ihfft(t);
+/// auto t = torch::fft::ihfft(T);
 /// assert(t.is_complex() && T.numel() == 128 / 2 + 1);
 /// ```
 inline Tensor ihfft(const Tensor& self,
@@ -226,6 +226,78 @@ inline Tensor ihfft(const Tensor& self,
   return torch::fft_ihfft(self, n, dim, norm);
 }
 
+/// Computes the 2-dimensional FFT of a Hermitian symmetric input signal.
+///
+/// The input is a onesided representation of the Hermitian symmetric time
+/// domain signal. See https://pytorch.org/docs/master/fft.html#torch.fft.hfft2.
+///
+/// Example:
+/// ```
+/// auto t = torch::randn({128, 65}, torch::kComplexDouble);
+/// auto T = torch::fft::hfft2(t, /*s=*/{128, 128});
+/// assert(T.is_floating_point() && T.numel() == 128 * 128);
+/// ```
+inline Tensor hfft2(const Tensor& self,
+                    c10::optional<IntArrayRef> s=c10::nullopt,
+                    IntArrayRef dim={-2, -1},
+                    c10::optional<c10::string_view> norm=c10::nullopt) {
+  return torch::fft_hfft2(self, s, dim, norm);
+}
+
+/// Computes the 2-dimensional IFFT of a real input signal.
+///
+/// The output is a onesided representation of the Hermitian symmetric time
+/// domain signal. See https://pytorch.org/docs/master/fft.html#torch.fft.ihfft2.
+///
+/// Example:
+/// ```
+/// auto T = torch::randn({128, 128}, torch::kDouble);
+/// auto t = torch::fft::hfft2(T);
+/// assert(t.is_complex() && t.size(1) == 65);
+/// ```
+inline Tensor ihfft2(const Tensor& self,
+                     c10::optional<IntArrayRef> s=c10::nullopt,
+                     IntArrayRef dim={-2, -1},
+                     c10::optional<c10::string_view> norm=c10::nullopt) {
+  return torch::fft_ihfft2(self, s, dim, norm);
+}
+
+/// Computes the N-dimensional FFT of a Hermitian symmetric input signal.
+///
+/// The input is a onesided representation of the Hermitian symmetric time
+/// domain signal. See https://pytorch.org/docs/master/fft.html#torch.fft.hfftn.
+///
+/// Example:
+/// ```
+/// auto t = torch::randn({128, 65}, torch::kComplexDouble);
+/// auto T = torch::fft::hfftn(t, /*s=*/{128, 128});
+/// assert(T.is_floating_point() && T.numel() == 128 * 128);
+/// ```
+inline Tensor hfftn(const Tensor& self,
+                    c10::optional<IntArrayRef> s=c10::nullopt,
+                    IntArrayRef dim={-2, -1},
+                    c10::optional<c10::string_view> norm=c10::nullopt) {
+  return torch::fft_hfftn(self, s, dim, norm);
+}
+
+/// Computes the N-dimensional IFFT of a real input signal.
+///
+/// The output is a onesided representation of the Hermitian symmetric time
+/// domain signal. See https://pytorch.org/docs/master/fft.html#torch.fft.ihfftn.
+///
+/// Example:
+/// ```
+/// auto T = torch::randn({128, 128}, torch::kDouble);
+/// auto t = torch::fft::hfft2(T);
+/// assert(t.is_complex() && t.size(1) == 65);
+/// ```
+inline Tensor ihfftn(const Tensor& self,
+                    c10::optional<IntArrayRef> s=c10::nullopt,
+                    IntArrayRef dim={-2, -1},
+                    c10::optional<c10::string_view> norm=c10::nullopt) {
+  return torch::fft_ihfftn(self, s, dim, norm);
+}
+
 /// Computes the discrete Fourier Transform sample frequencies for a signal of size n.
 ///
 /// See https://pytorch.org/docs/master/fft.html#torch.fft.fftfreq

torch/fft/__init__.py

@@ -874,7 +874,7 @@ Example:
 ihfft2 = _add_docstr(_fft.fft_ihfft2, r"""
 ihfft2(input, s=None, dim=(-2, -1), norm=None, *, out=None) -> Tensor
 
-Computes the N-dimensional inverse discrete Fourier transform of real
+Computes the 2-dimensional inverse discrete Fourier transform of real
 :attr:`input`. Equivalent to :func:`~torch.fft.ihfftn` but transforms only the
 two last dimensions by default.