code: 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],
      before_op: Operation,
  ) -> 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, before_op=before_op
      )
  
      # 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,), before_op=before_op
          )
          if input_imaginary:
              input_imaginary = mb.expand_dims(
                  x=input_imaginary, axes=(0,), before_op=before_op
              )
  
      is_onesided = 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
      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, before_op=before_op
          )
          if not window:
              window = mb.fill(
                  shape=(win_length.val,),
                  value=1.0,
                  before_op=before_op,
              )
          right = mb.fill(
              shape=(n_right,), value=0.0, before_op=before_op
          )
  
          # concatenate
          window = mb.concat(
              values=(left, window, right), axis=0, before_op=before_op
          )
  
      # apply time window
      if window:
          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)
  
      # conv with DFT kernel across the input signal
      sin_base = mb.sub(x=0.0, 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
      )
      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,
      )
  
      if input_imaginary:
          signal_imaginary = mb.expand_dims(
              x=input_imaginary, axes=(1,), 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,
          )
  
      # 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,
              before_op=before_op,
          )
          imag_result = mb.add(
              x=sin_windows_real,
              y=cos_windows_imag,
              before_op=before_op,
          )
      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,), before_op=before_op
          )
          imag_result = mb.squeeze(
              x=imag_result, axes=(0,), before_op=before_op
          )
  
      if normalized and normalized.val:
          divisor = mb.sqrt(
              x=mb.cast(x=n_fft, dtype="fp32", before_op=before_op),
              before_op=before_op,
          )
          real_result = mb.real_div(
              x=real_result, y=divisor, before_op=before_op
          )
          imag_result = mb.real_div(
              x=imag_result, y=divisor, before_op=before_op
          )
  
      return real_result, imag_result
  
imports: import functools
  from typing import Callable, Dict, Optional, Tuple
  import numpy
  from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.mil import types
  from coremltools.converters.mil.mil.operation import Operation
  from coremltools.converters.mil.mil.ops.defs.complex_dialect_ops import fft_canonicalize_length_dim, fft_canonicalize_shapes_dims
  from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
  from coremltools.converters.mil.mil.passes.helper import block_context_manager
  from coremltools.converters.mil.mil.passes.pass_registry import register_pass
  from coremltools.converters.mil.mil.var import ComplexVar, Var
label_desc: Use image node_id 6 for calling a global function with example usage: _stft(input_real, input_imaginary, n_fft, hop_length, win_length, window, normalized, onesided, before_op) and returns: real_result,  imag_result