More fixes

Author: Alejandro Gaston Alvarez Franceschi

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

@@ -9540,7 +9540,7 @@ def forward(self, x):
             [False, True], # center
             [False, True], # normalized
             [None, False, True], # onesided
-            [None, 60], # length
+            [None, 30, 40], # length
             [False, True], # return_complex
         )
     )
@@ -9582,6 +9582,14 @@ def forward(self, x):
                     backend=backend,
                     compute_unit=compute_unit
                 )
+        elif length is not None and return_complex is True:
+            with pytest.raises(ValueError, match="New var type `<class 'coremltools.converters.mil.mil.types.type_tensor.tensor.<locals>.tensor'>` not a subtype of existing var type `<class 'coremltools.converters.mil.mil.types.type_tensor.tensor.<locals>.tensor'>`"):
+                TorchBaseTest.run_compare_torch(
+                    input_shape,
+                    ISTFTModel(),
+                    backend=backend,
+                    compute_unit=compute_unit
+                )
         else:
             TorchBaseTest.run_compare_torch(
                 input_shape,

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

@@ -479,6 +479,8 @@ def _istft(
     window_mtx = mb.stack(values=[window_square] * n_frames, axis=0, before_op=before_op)
     window_mtx = mb.expand_dims(x=window_mtx, axes=(0,), before_op=before_op)
     window_envelope = _overlap_add(x=window_mtx, n_fft=n_fft, hop_length=hop_length, before_op=before_op)
+
+    # After this operation if it didn't have any channels dimention it adds one
     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
@@ -487,12 +489,8 @@ def _istft(
             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:
-                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:
-                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,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)
 
     return real_result, imag_result