code: def rolling_window(a: ndarray, window: int, hop_length: int):
      """return a stack of overlap subsequence of an array.
      ``return jnp.stack( [a[0:10], a[5:15], a[10:20],...], axis=0)``
      Source: https://github.com/google/jax/issues/3171
      Args:
        a (ndarray): input array of shape `[L, ...]`
        window (int): length of each subarray (window).
        hop_length (int): distance between neighbouring windows.
      """
  
      idx = (
          jnp.arange(window)[:, None]
          + jnp.arange((len(a) - window) // hop_length + 1)[None, :]
          * hop_length
      )
      return a[idx]
  
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 1 for calling a global function with example usage: rolling_window(a, window, hop_length) and returns: a