Add test

Author: Alejandro Gaston Alvarez Franceschi

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -9526,6 +9526,54 @@ def forward(self, x):
             compute_unit=compute_unit
         )
 
+class TestISTFT(TorchBaseTest):
+    @pytest.mark.slow
+    @pytest.mark.parametrize(
+        "compute_unit, backend, input_shape, complex, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, length",
+        itertools.product(
+            compute_units,
+            backends,
+            [(1, 32, 9), (32, 9), (3, 32, 9)], # input shape
+            [False, True], # complex
+            [16], # n_fft
+            [None, 4, 5], # hop_length
+            [None, 16, 9], # win_length
+            [None, torch.hann_window], # window
+            [None, False, True], # center
+            ["constant", "reflect", "replicate"], # pad mode
+            [False, True], # normalized
+            [None, False, True], # onesided
+            [None, 60], # length
+        )
+    )
+    def test_istft(self, compute_unit, backend, input_shape, complex, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided):
+        if complex and onesided:
+            pytest.skip("Onesided stft not possible for complex inputs")
+
+        class ISTFTModel(torch.nn.Module):
+            def forward(self, x):
+                applied_window = window(win_length) if window and win_length else None
+                x = torch.complex(x, x)
+                x = torch.istft(
+                    x,
+                    n_fft=n_fft,
+                    hop_length=hop_length,
+                    win_length=win_length,
+                    window=applied_window,
+                    center=center,
+                    normalized=normalized,
+                    onesided=onesided,
+                    length=length,
+                    return_complex=True)
+                x = torch.stack([torch.real(x), torch.imag(x)], dim=0)
+                return x
+
+        TorchBaseTest.run_compare_torch(
+            input_shape,
+            ISTFTModel(),
+            backend=backend,
+            compute_unit=compute_unit
+        )
 
 if _HAS_TORCH_AUDIO: