code: # Copyright 2020 NVIDIA. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import numpy as np
  from numba import jit
  
  
  def phase_vocoder(D: np.ndarray, rate: float, phi_advance: np.ndarray, scale_buffer: np.ndarray):
      """
      Optimized implementation of phase vocoder from Librosa.
  
      Reference implementation:
          - https://librosa.github.io/librosa/generated/librosa.core.phase_vocoder.html
  
      Args:
          D: Complex spectograms of shape [d, t, complex=2].
          rate: Speed rate, must be float greater than 0.
          phi_advance: Precomputed phase advance buffer array of length [n_fft + 1]
          scale_buffer: Precomputed numpy buffer array of length [n_fft + 1]
  
      Returns:
          Complex64 ndarray of shape [d, t / rate, complex=2]
      """
      time_steps = np.arange(0, D.shape[1], rate, dtype=np.float)
  
      # Create an empty output array
      d_stretch = np.zeros((D.shape[0], len(time_steps)), D.dtype, order='F')
  
      # Phase accumulator; initialize to the first sample
      phase_acc = np.angle(D[:, 0])
  
      # Pad 0 columns to simplify boundary logic
      D = np.pad(D, [(0, 0), (0, 2)], mode='constant')
  
      d_stretch = _phase_vocoder_kernel(D, time_steps, phi_advance, d_stretch, phase_acc, scale_buffer)
  
      return d_stretch
  
  
  @jit(nopython=True, nogil=True)
  def _phase_vocoder_kernel(D, time_steps, phi_advance, d_stretch, phase_acc, scale_buffer):
      """
      Numba optimized kernel to compute the phase vocoder step.
  
      Args:
          D: Complex spectograms of shape [d, t, complex=2].
          rate: Speed rate, must be float greater than 0.
          time_steps: Numpy ndarray of linearly spaced time steps, shape = [t]
          phi_advance: Precomputed phase advance buffer array of length [n_fft + 1]
          d_stretch: Output complex matrix of shape [d, t / rate, complex=2]
          phase_acc: Phase accumulator initialized to first sample of shape [d, complex=2]
          scale_buffer: Precomputed numpy buffer array of length [n_fft + 1]
  
      Returns:
          Complex64 ndarray of shape [d, t / rate, complex=2]
      """
      two_pi = 2.0 * np.pi
  
      for (t, step) in enumerate(time_steps):
          columns = D[:, int(step) : int(step + 2)]
          columns_0 = columns[:, 0]
          columns_1 = columns[:, 1]
  
          # Weighting for linear magnitude interpolation
          alpha = np.mod(step, 1.0)
          mag = (1.0 - alpha) * np.abs(columns_0) + alpha * np.abs(columns_1)
  
          # Store to output array
          d_stretch[:, t] = mag * np.exp(1.0j * phase_acc)
  
          # Compute phase advance
          dphase = np.angle(columns_1) - np.angle(columns_0) - phi_advance
  
          # Wrap to -pi:pi range
          scale = dphase / two_pi
          np.round(scale, 0, scale_buffer)
  
          dphase = dphase - two_pi * scale_buffer
  
          # Accumulate phase
          phase_acc += phi_advance + dphase
  
      return d_stretch
  
apis: numpy.abs numpy.exp numpy.mod numpy.arange numpy.angle numpy.round numpy.pad