Fixes

Author: Alejandro Gaston Alvarez Franceschi

coremltools/converters/mil/frontend/torch/ops.py

@@ -6362,6 +6362,7 @@ def stft(context, node):
     Lowers torch.stft with the dialect op `complex_stft` from complex_dialect_ops.py
     """
     input_data, n_fft, hop_length, win_length, window, normalized, onesided, _ = _get_inputs(context, node, min_expected=2)
+
     if types.is_complex(input_data.dtype):
         onesided = False # pytorch defaults onesided to False for complex inputs
     stft_res = mb.complex_stft(
@@ -6371,9 +6372,32 @@ def stft(context, node):
         win_length=win_length,
         window=window,
         normalized=normalized,
-        onesided=onesided)
+        onesided=onesided
+    )
     context.add(stft_res, node.name)
 
+@register_torch_op
+def istft(context, node):
+    """
+    Lowers torch.istft with the dialect op `complex_istft` from complex_dialect_ops.py
+    """
+    input_data, n_fft, hop_length, win_length, window, center, normalized, onesided, length, _ = _get_inputs(context, node, min_expected=2)
+
+    if types.is_complex(input_data.dtype):
+        onesided = False # pytorch defaults onesided to False for complex inputs
+    istft_res = mb.complex_istft(
+        input=input_data,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        normalized=normalized,
+        onesided=onesided,
+        length=length,
+    )
+    context.add(istft_res, node.name)
+
 @register_torch_op(torch_alias=["torchvision::nms"])
 def torchvision_nms(context, node):
     inputs = _get_inputs(context, node, expected=3)

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

@@ -9589,12 +9589,11 @@ def forward(self, x):
 
     @pytest.mark.slow
     @pytest.mark.parametrize(
-        "compute_unit, backend, input_shape, n_fft, hop_length, win_length, window, center, normalized, onesided, length, return_complex",
+        "compute_unit, backend, input_shape, hop_length, win_length, window, center, normalized, onesided, length, return_complex",
         itertools.product(
             compute_units,
             backends,
             [(1, 32, 9), (32, 9), (3, 32, 9)], # input shape
-            [16], # n_fft
             [None, 4, 5], # hop_length
             [None, 16, 9], # win_length
             [None, torch.hann_window], # window
@@ -9605,10 +9604,12 @@ def forward(self, x):
             [False, True], # return_complex
         )
     )
-    def test_istft(self, compute_unit, backend, input_shape, n_fft, hop_length, win_length, window, center, normalized, onesided, length, return_complex):
+    def test_istft(self, compute_unit, backend, input_shape, hop_length, win_length, window, center, normalized, onesided, length, return_complex):
         if return_complex and onesided:
             pytest.skip("Complex output is incompatible with onesided")
 
+        n_fft = input_shape[1]
+
         class ISTFTModel(torch.nn.Module):
             def forward(self, x):
                 applied_window = window(win_length) if window and win_length else None

coremltools/converters/mil/mil/ops/defs/complex_dialect_ops.py

@@ -893,7 +893,6 @@ class complex_istft(Operation):
 
     Attributes
     ----------
-    V: complex64
     T: fp32, complex64
 
     References
@@ -902,44 +901,40 @@ class complex_istft(Operation):
     """
 
     input_spec = InputSpec(
-        input=TensorInputType(type_domain="V"),
+        input=TensorInputType(type_domain=types.complex),
         n_fft=TensorInputType(const=True, type_domain=types.int32),
         hop_length=TensorInputType(const=True, optional=True, type_domain=types.int32),
         win_length=TensorInputType(const=True, optional=True, type_domain=types.int32),
         window=TensorInputType(const=True, optional=True, type_domain=types.fp32),
-        normalized=TensorInputType(const=True, optional=True, type_domain=types.bool),
+        center=TensorInputType(const=True, type_domain=types.bool),
+        normalized=TensorInputType(const=True, optional=False, type_domain=types.bool),
         onesided=TensorInputType(const=True, optional=True, type_domain=types.bool),
         length=TensorInputType(const=True, optional=True, type_domain=types.int32),
+        return_complex=TensorInputType(const=True, optional=True, type_domain=types.bool),
     )
 
-    type_domains = {
-        "V": types.complex64,
-    }
-
     def default_inputs(self):
         return DefaultInputs(
             hop_length = None,
             win_length = None,
             window = None,
             normalized = False,
             onesided = True,
-            length = None
+            length = None,
+            return_complex = True,
         )
 
     def type_inference(self):
-        output_type = (types.fp32)
-        output_shape = []
+        output_type = (types.complex64) if self.return_complex else (types.fp32)
 
-        # add back rank if needed
-        if self.input.rank == 2:
-            output_shape += [self.input.shape[0]]
+        # add batch size if given
+        output_shape = [self.input.shape[0] if self.input.rank == 3 else 1]
 
         if self.length:
             output_shape += [self.length]
-            return types.tensor(output_type, tuple(output_shape))
+        else:
+            n_frames = self.input.shape[-1]
+            output_shape += [self.n_fft.val + self.hop_length.val * (n_frames - 1)]
 
-        n_frames = self.input.shape[-1]
-        output_shape = self.n_fft.val + self.hop_length.val * (n_frames - 1)
 
         return types.tensor(output_type, tuple(output_shape))
-

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

@@ -397,6 +397,7 @@ def _istft(
     hop_length: Optional[Var],
     win_length: Optional[Var],
     window: Optional[Var],
+    center: Optional[Var],
     normalized: Optional[Var],
     onesided: Optional[Var],
     length: Optional[Var],
@@ -435,12 +436,10 @@ def _istft(
         cos_base = mb.mul(x=window, y=cos_base, before_op=before_op)
         sin_base = mb.mul(x=window, y=sin_base, before_op=before_op)
 
-    cos_base = mb.expand_dims(x=cos_base, axes=(1,), before_op=before_op)
-    sin_base = mb.expand_dims(x=sin_base, axes=(1,), before_op=before_op)
     hop_size = mb.expand_dims(x=hop_length, axes=(0,), before_op=before_op)
 
-    signal_real = mb.expand_dims(x=input_real, axes=(1,), before_op=before_op)
-    signal_imaginary = mb.expand_dims(x=input_imaginary, axes=(1,), before_op=before_op)
+    signal_real = input_real
+    signal_imaginary = input_imaginary
 
     # De-normalized signal before applying the IFT
     if normalized and normalized.val:
@@ -455,15 +454,16 @@ def _istft(
     # So using the definition in stft function, we get:
     #   real(x[n]) = Σ(a[k]*cos(2π*k/K*n)+b[k]*sin(2π*k/K*n))
     #   imag(x[n]) = Σ(b[k]*cos(2π*k/K*n)-a[k]*sin(2π*k/K*n))
-    cos_windows_real = mb.conv(x=signal_real, weight=cos_base, strides=hop_size, pad_type='valid', before_op=before_op)
-    sin_windows_real = mb.conv(x=signal_real, weight=sin_base, strides=hop_size, pad_type='valid', before_op=before_op)
-    cos_windows_imag = mb.conv(x=signal_imaginary, weight=cos_base, strides=hop_size, pad_type='valid', before_op=before_op)
-    sin_windows_imag = mb.conv(x=signal_imaginary, weight=sin_base, strides=hop_size, pad_type='valid', before_op=before_op)
+    cos_windows_real = mb.matmul(x=signal_real, y=cos_base, transpose_x=True, before_op=before_op)
+    sin_windows_real = mb.matmul(x=signal_real, y=sin_base, transpose_x=True, before_op=before_op)
+    cos_windows_imag = mb.matmul(x=signal_imaginary, y=cos_base, transpose_x=True, before_op=before_op)
+    sin_windows_imag = mb.matmul(x=signal_imaginary, y=sin_base, transpose_x=True, before_op=before_op)
 
     real_result = mb.add(x=cos_windows_real, y=sin_windows_imag, before_op=before_op)
     imag_result = mb.sub(x=cos_windows_imag, y=sin_windows_real, before_op=before_op)
 
     # Divide by N
+    n_fft = mb.cast(x=n_fft, dtype="fp32", before_op=before_op)
     real_result = mb.real_div(x=real_result, y=n_fft, before_op=before_op)
     imag_result = mb.real_div(x=imag_result, y=n_fft, before_op=before_op)
 
@@ -472,9 +472,9 @@ def _istft(
     imag_result = _overlap_add(x=imag_result, n_fft=n_fft, hop_length=hop_length, before_op=before_op)
 
     # Normalize by the window square
-    n_frames = mb.shape(x=real_result, before_op=before_op)[1]
     window_square = mb.mul(x=window, y=window, before_op=before_op)
-    window_mtx = mb.stack(values=[window_square] * n_frames, axis=1)
+    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)
     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)
@@ -502,17 +502,27 @@ def _overlap_add(
     """
     input_shape = mb.shape(x=x, before_op=before_op)
     channels = input_shape.val[0]
-    n_frames = input_shape.val[2]
+    n_frames = input_shape.val[1]
+
+    # Create empty output with final shape
+    output = mb.fill(shape=(channels, int(n_fft.val + hop_length.val * (n_frames - 1))), value=0., before_op=before_op)
 
-    output = mb.fill(shape=(channels, n_fft.val + hop_length.val * (n_frames - 1)), value=0., before_op=before_op)
-    signal_frames = mb.split(x=x, num_splits=n_frames, axis=2, before_op=before_op)
+    # Create an index used later on overlap add
+    n_fft = mb.cast(x=n_fft, dtype="int32", before_op=before_op)
     local_idx = mb.range_1d(start=0, end=n_fft, step=1, before_op=before_op)
 
+    # Split data into frames and iterate
+    signal_frames = mb.split(x=x, num_splits=n_frames, axis=1, before_op=before_op)
+
     for frame_num, frame in enumerate(signal_frames):
+        frame = mb.squeeze(x=frame, axes=[1], before_op=before_op)
+
+        # Create index to align data frames
         global_idx = mb.add(x=local_idx , y=frame_num*hop_length.val, before_op=before_op)
-        global_idx = mb.expand_dims(x=global_idx, axes=(0,), before_op=before_op)
-        global_idx = mb.stack(values=[global_idx] * channels, axis=0)
-        output = mb.scatter_nd(data=output, indices=global_idx, updates=frame, before_op=before_op)
+        global_idx = mb.stack(values=[global_idx] * channels, axis=0, before_op=before_op)
+
+        # Add data frame
+        output = mb.scatter_along_axis(data=output, indices=global_idx, updates=frame, axis=1, mode="add", before_op=before_op)
 
     return output
 
@@ -748,19 +758,18 @@ def _lower_complex_stft(op: Operation):
 
 @LowerComplex.register_lower_func(op_type="complex_istft")
 def _lower_complex_istft(op: Operation):
-    is_complex = types.is_complex(op.input.dtype)
 
     # check parameters for validity
-    if is_complex:
-        raise ValueError("Only complex inputs are allowed")
+    if not types.is_complex(op.input.dtype):
+        raise TypeError("Input type must be complex")
     if op.win_length and op.win_length.val > op.n_fft.val:
         raise ValueError("Window length must be less than or equal to n_fft")
     if op.return_complex and op.onesided and op.onesided.val:
         raise ValueError("Complex output is not compatible with onesided")
 
     real, imag = _istft(
         op.input.real, op.input.imag,
-        op.n_fft, op.hop_length, op.win_length, op.window, op.normalized, op.onesided, op.length, before_op=op)
+        op.n_fft, op.hop_length, op.win_length, op.window, op.center, op.normalized, op.onesided, op.length, before_op=op)
 
     if op.return_complex:
         return _wrap_complex_output(op.outputs[0], real, imag)