def _stft(
    input_real: Var,
    input_imaginary: Optional[Var],
    n_fft: Var,
    hop_length: Optional[Var],
    win_length: Optional[Var],
    window: Optional[Var],
    normalized: Optional[Var],
    onesided: Optional[Var],
) -> Tuple[Var, Var]:
    """
    We can write STFT in terms of convolutions with a DFT kernel.
    At the end:
        * The real part output is: cos_base * input_real + sin_base * input_imag
        * The imaginary part output is: - (sin_base * input_real - cos_base * input_imag)
    Adapted from: https://github.com/adobe-research/convmelspec/blob/main/convmelspec/mil.py
    """
    hop_length = hop_length or mb.floor_div(x=n_fft, y=4)

    # input should always be 2D
    should_increase_rank = input_real.rank == 1
    if should_increase_rank:
        input_real = mb.expand_dims(x=input_real, axes=(0,))
        if input_imaginary:
            input_imaginary = mb.expand_dims(x=input_imaginary, axes=(0,))

    is_onesided = onesided and onesided.val
    cos_base, sin_base = _calculate_dft_matrix(n_fft, onesided=is_onesided)

    # create a window of centered 1s of the requested size
    if win_length:
        n_left = (n_fft.val - win_length.val) // 2
        n_right = n_fft.val - win_length.val - n_left

        left = mb.fill(shape=(n_left,), value=0.0)
        if not window:
            window = mb.fill(shape=(win_length.val,), value=1.0)
        right = mb.fill(shape=(n_right,), value=0.0)

        # concatenate
        window = mb.concat(values=(left, window, right), axis=0)

    # apply time window
    if window:
        cos_base = mb.mul(x=window, y=cos_base)
        sin_base = mb.mul(x=window, y=sin_base)

    # conv with DFT kernel across the input signal
    sin_base = mb.sub(x=0.0, y=sin_base)
    cos_base = mb.expand_dims(x=cos_base, axes=(1,))
    sin_base = mb.expand_dims(x=sin_base, axes=(1,))
    hop_size = mb.expand_dims(x=hop_length, axes=(0,))

    signal_real = mb.expand_dims(x=input_real, axes=(1,))
    cos_windows_real = mb.conv(x=signal_real, weight=cos_base, strides=hop_size, pad_type="valid")
    sin_windows_real = mb.conv(x=signal_real, weight=sin_base, strides=hop_size, pad_type="valid")

    if input_imaginary:
        signal_imaginary = mb.expand_dims(x=input_imaginary, axes=(1,))
        cos_windows_imag = mb.conv(
            x=signal_imaginary, weight=cos_base, strides=hop_size, pad_type="valid"
        )
        sin_windows_imag = mb.conv(
            x=signal_imaginary, weight=sin_base, strides=hop_size, pad_type="valid"
        )

    # add everything together
    if input_imaginary:
        # sin base is already negative so subtract
        real_result = mb.sub(x=cos_windows_real, y=sin_windows_imag)
        imag_result = mb.add(x=sin_windows_real, y=cos_windows_imag)
    else:
        real_result = cos_windows_real
        imag_result = sin_windows_real

    # reduce the rank of the output
    if should_increase_rank:
        real_result = mb.squeeze(x=real_result, axes=(0,))
        imag_result = mb.squeeze(x=imag_result, axes=(0,))

    if normalized and normalized.val:
        divisor = mb.sqrt(x=mb.cast(x=n_fft, dtype="fp32"))
        real_result = mb.real_div(x=real_result, y=divisor)
        imag_result = mb.real_div(x=imag_result, y=divisor)

    return real_result, imag_result

    We can write STFT in terms of convolutions with a DFT kernel.
    At the end:
        * The real part output is: cos_base * input_real + sin_base * input_imag
        * The imaginary part output is: - (sin_base * input_real - cos_base * input_imag)
    Adapted from: https://github.com/adobe-research/convmelspec/blob/main/convmelspec/mil.py
    
coremltools/converters/mil/mil/passes/defs/lower_complex_dialect_ops.py