Add autograd test to T.TimeStretch (and F.phase_vocoder)

Author: mthrok

test/torchaudio_unittest/transforms/autograd_test_impl.py

@@ -1,4 +1,5 @@
 from typing import List
+import unittest
 
 from parameterized import parameterized
 import torch
@@ -37,8 +38,12 @@ def assert_grad(
 
         inputs_ = []
         for i in inputs:
-            i.requires_grad = True
-            inputs_.append(i.to(dtype=torch.float64, device=self.device))
+            if torch.is_tensor(i):
+                i = i.to(
+                    dtype=torch.cdouble if i.is_complex() else torch.double,
+                    device=self.device)
+                i.requires_grad = True
+            inputs_.append(i)
         assert gradcheck(transform, inputs_)
         assert gradgradcheck(transform, inputs_, nondet_tol=nondet_tol)
 
@@ -123,3 +128,52 @@ def test_spectral_centroid(self):
         transform = T.SpectralCentroid(sample_rate=sample_rate)
         waveform = get_whitenoise(sample_rate=sample_rate, duration=0.05, n_channels=2)
         self.assert_grad(transform, [waveform], nondet_tol=1e-10)
+
+    @unittest.expectedFailure
+    def test_timestretch_zeros_fail(self):
+        """Test that ``T.TimeStretch`` fails gradcheck at 0
+
+        This is because ``F.phase_vocoder`` converts data from cartesian to polar coordinate,
+        which performs ``atan2(img, real)``, and gradient is not defined at 0.
+        """
+        n_fft = 16
+        transform = T.TimeStretch(n_freq=n_fft // 2 + 1, fixed_rate=0.99)
+        waveform = torch.zeros(2, 40)
+        spectrogram = get_spectrogram(waveform, n_fft=n_fft, power=None)
+        self.assert_grad(transform, [spectrogram])
+
+    @nested_params(
+        [0.7, 0.8, 0.9, 1.0, 1.3],
+        [False, True],
+    )
+    def test_timestretch(self, rate, test_pseudo_complex):
+        """Verify that ``T.TimeStretch`` does not fail if it's not too close to 0
+
+        ``T.TimeStrech`` is not differentiable around 0, so this test checks the differentiability
+        for cases where input is not zero, and different configurations of `TimeStretch`.
+
+        Ideally, we should be testing on Spectrogram of random waveform but it is hard to control
+        the values around zeros.
+        """
+        n_fft = 16
+        transform = T.TimeStretch(n_freq=n_fft // 2 + 1, fixed_rate=rate)
+        waveform = torch.zeros(2, 40)
+        spectrogram = get_spectrogram(waveform, n_fft=n_fft, power=None)
+
+        # Epsilon values tried
+        #
+        # Note:
+        #   This is not experimental and comprehensive.
+        #   The result also depends on ``n_fft``.
+        #
+        #                 CPU / CUDA
+        # * 1e-1           ok / ok
+        # * 1e-2           ok / ok
+        # * 1e-3           ok / ok
+        # * 1e-3 + 1e-3j   ok / ok
+        # * 1e-4           ok / NG
+
+        spectrogram += 1e-3
+        if test_pseudo_complex:
+            spectrogram = torch.view_as_real(spectrogram)
+        self.assert_grad(transform, [spectrogram])