code: def batched_stft(
      y: ndarray,
      n_fft: int,
      hop_length: int,
      win_length: int,
      window: str,
      center: bool = True,
      pad_mode: str = "reflect",
  ):
      """Batched version of ``stft`` function.
      TN => FTN
      """
  
      assert len(y.shape) >= 2
      if window == "hann":
          fft_window = jnp.hanning(win_length + 1)[:-1]
      else:
          raise RuntimeError(
              f"{window} window function is not supported!"
          )
      pad_len = (n_fft - win_length) // 2
      if pad_len > 0:
          fft_window = jnp.pad(
              fft_window, (pad_len, pad_len), mode="constant"
          )
          win_length = n_fft
      else:
          fft_window = fft_window
      if center:
          pad_width = ((n_fft // 2, n_fft // 2),) + ((0, 0),) * (
              len(y.shape) - 1
          )
          y = jnp.pad(y, pad_width, mode=pad_mode)
  
      # jax does not support ``np.lib.stride_tricks.as_strided`` function
      # see https://github.com/google/jax/issues/3171 for comments.
      y_frames = rolling_window(y, n_fft, hop_length)
      fft_window = jnp.reshape(
          fft_window, (-1,) + (1,) * (len(y.shape))
      )
      y_frames = y_frames * fft_window
      stft_matrix = jnp.fft.fft(y_frames, axis=0)
      d = int(1 + n_fft // 2)
      return stft_matrix[:d]
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/dsp.py
imports: from functools import partial
  from typing import Optional
  import jax
  import jax.numpy
  import librosa
  from einops import rearrange
  from jax.numpy import ndarray
label_desc: Use image node_id 3 for calling a global function with example usage: batched_stft(y, n_fft, hop_length, win_length, window, center, pad_mode) and returns: stft_matrix