code: def stft(
      data,
      n_fft,
      hop_length,
      win_length,
      window,
      normalized,
      onesided,
      output_shape,
  ):
      """
      The STFT computes the Fourier transform of short overlapping windows of the input.
      This gives frequency components of the signal as they change over time.
      Parameters
      ----------
      data : relay.Expr
          Either a 1-D tensor or a 2-D batch tensor.
      n_fft : int
          The size of Fourier transform
      hop_length : int
          The distance between neighboring sliding window frames
      win_length : int
          The size of window frame and STFT filter
      window : relay.Expr
          A 1-D tensor window frame
      normalized : bool
          Whether to return the normalized STFT results
      onesided : bool
          Whether to return onesided result or fill with conjugate symmetry
      Returns
      -------
      output : relay.Expr
          Tensor containing the STFT result
      Examples
      --------
      .. code-block:: python
  
          data = [1, 2, 3, 4, 5, 6]
          window = [4, 3, 2]
          [n_fft, hop_length, win_length, normalized, onesided] = [3, 3, 3, False, True]
          relay.stft(data, n_fft, hop_length, win_length, window, normalized, onesided)
          -> [[[15.0000,  0.0000], [34.0000,  0.0000]], [[ 4.5000,  0.8660], [ 1.0000, -1.7321]]]
      """
  
      def gen_ir(
          data_ptr,
          n_fft,
          hop_length,
          win_length,
          window_ptr,
          normalized,
          onesided,
          output_ptr,
          loop_kind,
      ):
          ib = tir.ir_builder.create()
          data = ib.buffer_ptr(data_ptr)
          window = ib.buffer_ptr(window_ptr)
          output = ib.buffer_ptr(output_ptr)
          # https://librosa.org/doc/0.7.2/_modules/librosa/core/spectrum.html#stft
          with ib.for_range(
              0,
              output_ptr.shape[0] * output_ptr.shape[1],
              kind="parallel",
          ) as batch_row:
              with ib.for_range(
                  0, output_ptr.shape[2], kind=loop_kind
              ) as col:
                  batch = ib.allocate(
                      "int32", (1), name="batch", scope="local"
                  )
                  row = ib.allocate(
                      "int32", (1), name="row", scope="local"
                  )
                  batch = tir.floordiv(batch_row, output_ptr.shape[1])
                  row = tir.floormod(batch_row, output_ptr.shape[1])
                  output[batch, row, col, 0] = tir.Cast(
                      data_ptr.dtype, 0
                  )
                  output[batch, row, col, 1] = tir.Cast(
                      data_ptr.dtype, 0
                  )
                  with ib.for_range(0, win_length) as wlen:
                      output[batch, row, col, 0] += (
                          window[wlen]
                          * data[batch, col * hop_length + wlen]
                          * tir.cos(2 * pi * row * wlen / win_length)
                      )
                      output[batch, row, col, 1] -= (
                          window[wlen]
                          * data[batch, col * hop_length + wlen]
                          * tir.sin(2 * pi * row * wlen / win_length)
                      )
                  with ib.if_scope(normalized):
                      output[batch, row, col, 0] /= tir.sqrt(
                          tir.const(n_fft, "float32")
                      )
                      output[batch, row, col, 1] /= tir.sqrt(
                          tir.const(n_fft, "float32")
                      )
  
          return ib.get()
  
      output_buf = tir.decl_buffer(
          output_shape, data.dtype, "output_buf"
      )
      loop_kind = "vectorize"
      if isinstance(output_shape[2], tir.expr.SizeVar):  # any_dim
          loop_kind = "serial"
  
      return te.extern(
          output_shape,
          [data, window],
          lambda ins, outs: gen_ir(
              ins[0],
              n_fft,
              hop_length,
              win_length,
              ins[1],
              normalized,
              onesided,
              outs[0],
              loop_kind,
          ),
          dtype=[data.dtype],
          out_buffers=[output_buf],
          name="stft_cpu",
          tag="stft_cpu",
      )
  
imports: from math import pi
  from tvm import te, tir
label_desc: Use image node_id 1 for calling a global function with example usage: stft(data, n_fft, hop_length, win_length, window, normalized, onesided, output_shape) and returns: te, ib