Address review comments

peterbell10 · peterbell10 · commit 85da7fef5439 · 2020-09-15T12:51:56.000+01:00
diff --git a/aten/src/ATen/native/SpectralOps.cpp b/aten/src/ATen/native/SpectralOps.cpp
@@ -193,6 +193,14 @@ static Stream& write_opt(Stream& SS, const optional<T>& value) {
   return SS;
 }
 
+/* Short-time Fourier Transform, for signal analysis.
+ *
+ * This is modelled after librosa but with support for complex time-domain
+ * signals and complex windows.
+ *
+ * NOTE: librosa's center and pad_mode arguments are currently only implemented
+ * in python because it uses torch.nn.functional.pad which is python-only.
+ */
 Tensor stft(const Tensor& self, const int64_t n_fft, const optional<int64_t> hop_lengthOpt,
             const optional<int64_t> win_lengthOpt, const Tensor& window,
             const bool normalized, const optional<bool> onesidedOpt,
@@ -306,6 +314,11 @@ Tensor stft(const Tensor& self, const int64_t n_fft, const optional<int64_t> hop
   }
 }
 
+/* Inverse Short-time Fourier Transform
+ *
+ * This is modelled after librosa but with support for complex time-domain
+ * signals and complex windows.
+ */
 Tensor istft(const Tensor& self, const int64_t n_fft, const optional<int64_t> hop_lengthOpt,
              const optional<int64_t> win_lengthOpt, const Tensor& window,
              const bool center, const bool normalized, const c10::optional<bool> onesidedOpt,
diff --git a/test/test_spectral_ops.py b/test/test_spectral_ops.py
@@ -30,6 +30,10 @@ def _complex_stft(x, *args, **kwargs):
 
 
 def _hermitian_conj(x, dim):
+    """Returns the hermitian conjugate along a single dimension
+
+    H(x)[i] = conj(x[-i])
+    """
     out = torch.empty_like(x)
     mid = (x.size(dim) - 1) // 2
     idx = [slice(None)] * out.dim()
@@ -60,10 +64,17 @@ def _complex_istft(x, *args, **kwargs):
     x_antihermitian = (x - hconj) / 2
     istft_real = torch.istft(x_hermitian[slc], *args, **kwargs, onesided=True)
     istft_imag = torch.istft(-1j * x_antihermitian[slc], *args, **kwargs, onesided=True)
-    return istft_real + 1j * istft_imag
+    return torch.complex(istft_real, istft_imag)
+
+
+def _stft_reference(x, hop_length, window):
+    r"""Reference stft implementation
+
+    This doesn't implement all of torch.stft, only the STFT definition:
 
+    .. math:: X(m, \omega) = \sum_n x[n]w[n - m] e^{-jn\omega}
 
-def stft_definition(x, hop_length, window):
+    """
     n_fft = window.numel()
     X = torch.empty((n_fft, (x.numel() - n_fft + hop_length) // hop_length),
                     device=x.device, dtype=torch.cdouble)
@@ -482,7 +493,7 @@ def test_complex_stft_definition(self, device, dtype):
 
         for args in test_args:
             window = torch.randn(args[1], device=device, dtype=dtype)
-            expected = stft_definition(args[0], args[2], window)
+            expected = _stft_reference(args[0], args[2], window)
             actual = torch.stft(*args, window=window, center=False)
             self.assertEqual(actual, expected)
 
diff --git a/torch/functional.py b/torch/functional.py
@@ -399,6 +399,12 @@ def stft(input: Tensor, n_fft: int, hop_length: Optional[int] = None,
          return_complex: Optional[bool] = None) -> Tensor:
     r"""Short-time Fourier transform (STFT).
 
+    The STFT computes the Fourier transform of short overlapping windows of the
+    input. This giving frequency components of the signal as they change over
+    time. The interface of this function is modelled after librosa_.
+
+    .. _librosa: https://librosa.org/doc/latest/generated/librosa.stft.html
+
     Ignoring the optional batch dimension, this method computes the following
     expression:
 
@@ -452,10 +458,10 @@ def stft(input: Tensor, n_fft: int, hop_length: Optional[int] = None,
       dimension represents the real and imaginary components.
 
     Returns either a complex tensor of size :math:`(* \times N \times T)` if
-    :attr:`return_complex`, or a real tensor of size :math:`(* \times N \times
-    T \times 2)`. :math:`*` is the optional batch size of :attr:`input`,
-    :math:`N` is the number of frequencies where STFT is applied and :math:`T`
-    is the total number of frames used.
+    :attr:`return_complex` is true, or a real tensor of size :math:`(* \times N
+    \times T \times 2)`. Where :math:`*` is the optional batch size of
+    :attr:`input`, :math:`N` is the number of frequencies where STFT is applied
+    and :math:`T` is the total number of frames used.
 
     .. warning::
       This function changed signature at version 0.4.1. Calling with the
@@ -479,7 +485,7 @@ def stft(input: Tensor, n_fft: int, hop_length: Optional[int] = None,
              Default: ``False``
         onesided (bool, optional): controls whether to return half of results to
             avoid redundancy for real inputs.
-            Default: ``True`` for real input and window, ``False`` otherwise.
+            Default: ``True`` for real :attr:`input` and :attr:`window`, ``False`` otherwise.
         return_complex (bool, optional): whether to return a complex tensor, or
             a real tensor with an extra last dimension for the real and
             imaginary components.
