Fixes

Author: Alejandro Gaston Alvarez Franceschi

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

@@ -9597,7 +9597,7 @@ def forward(self, x):
             [16, 32], # n_fft
             [5, 9], # num_frames
             [None, 4, 5], # hop_length
-            [None, 16, 9], # win_length
+            [None, 10, 8], # win_length
             [None, torch.hann_window], # window
             [False, True], # center
             [False, True], # normalized
@@ -9610,6 +9610,9 @@ def test_istft(self, compute_unit, backend, channels, n_fft, num_frames, hop_len
         if return_complex and onesided:
             pytest.skip("Complex output is incompatible with onesided")
 
+        if hop_length is None and win_length is not None:
+            pytest.skip("If win_length is set then we must set hop_length and 0 < hop_length <= win_length")
+
         freq = n_fft//2+1 if onesided else n_fft
         input_shape = (channels, freq, num_frames) if channels else (freq, num_frames)
 
@@ -9628,24 +9631,13 @@ def forward(self, x):
                     length=length,
                     return_complex=return_complex)
                 if return_complex:
-                    x = torch.stack([torch.real(x), torch.imag(x)], dim=0)
-                return x
+                    return torch.stack([torch.real(x), torch.imag(x)], dim=0)
+                else:
+                    return torch.real(x)
 
-        if length is not None or center is False:
+        if win_length and center is False:
             # For some reason Pytorch raises an error https://github.com/pytorch/audio/issues/427#issuecomment-1829593033
-            with pytest.raises(
-                RuntimeError, match="istft\(.*\) window overlap add min: 1"
-            ):
-                TorchBaseTest.run_compare_torch(
-                    input_shape,
-                    ISTFTModel(),
-                    backend=backend,
-                    compute_unit=compute_unit
-                )
-        elif return_complex is False:
-            with pytest.raises(
-                ValueError, match="MIL doesn't support complex data as model's output"
-            ):
+            with pytest.raises(RuntimeError, match="istft\(.*\) window overlap add min: 1"):
                 TorchBaseTest.run_compare_torch(
                     input_shape,
                     ISTFTModel(),

coremltools/converters/mil/mil/passes/defs/lower_complex_dialect_ops.py

@@ -427,7 +427,7 @@ def _istft(
 
     expected_output_signal_len = n_fft.val + hop_length.val * (n_frames - 1)
 
-    is_onesided = onesided.val if onesided else fft_size != n_fft
+    is_onesided = True if fft_size != n_fft.val else onesided and onesided.val
     cos_base, sin_base = _calculate_dft_matrix(n_fft, onesided=is_onesided, before_op=before_op)
 
     # create a window of centered 1s of the requested size
@@ -481,20 +481,18 @@ def _istft(
     window_envelope = _overlap_add(x=window_mtx, n_fft=n_fft, hop_length=hop_length, before_op=before_op)
     real_result = mb.real_div(x=real_result, y=window_envelope, before_op=before_op)
     imag_result = mb.real_div(x=imag_result, y=window_envelope, before_op=before_op)
-
     # We need to adapt last dimension
     if length is not None:
         if length.val > expected_output_signal_len:
+            real_result = mb.pad(x=real_result, pad=(0, length.val - expected_output_signal_len), before_op=before_op)
+            imag_result = mb.pad(x=imag_result, pad=(0, length.val - expected_output_signal_len), before_op=before_op)
+        elif length.val < expected_output_signal_len:
             if channels:
-                right_pad = mb.fill(shape=(channels, length.val - expected_output_signal_len ), value=0., before_op=before_op)
+                real_result = mb.slice_by_size(x=real_result, begin=[0,0], size=[-1, length.val], before_op=before_op)
+                imag_result = mb.slice_by_size(x=imag_result, begin=[0,0], size=[-1, length.val], before_op=before_op)
             else:
-                right_pad = mb.fill(shape=(length.val - expected_output_signal_len,), value=0., before_op=before_op)
-
-            real_result = mb.stack(values=(real_result, right_pad), axis=1, before_op=before_op)
-            imag_result = mb.stack(values=(imag_result, right_pad), axis=1, before_op=before_op)
-        elif length.val < expected_output_signal_len:
-            real_result = mb.slice_by_size(x=real_result, begin=[0], size=[length.val], before_op=before_op)
-            imag_result = mb.slice_by_size(x=imag_result, begin=[0], size=[length.val], before_op=before_op)
+                real_result = mb.slice_by_size(x=real_result, begin=[0], size=[length.val], before_op=before_op)
+                imag_result = mb.slice_by_size(x=imag_result, begin=[0], size=[length.val], before_op=before_op)
 
     return real_result, imag_result