code: def rolling_window(
      x: jnp.ndarray,
      window: int,
      axis: int = 0,
      time_major: bool = True,
  ):
      """Stack the N=T-W+1 length W slices [0:W, 1:W+1, ..., T-W:T] from a tensor.
  
      Args:
        x: The tensor to select rolling slices from (along specified axis), with
          shape [..., T, ...]; i.e., T = x.shape[axis].
        window: The length (W) of the slices to select.
        axis: The axis to slice from (defaults to 0).
        time_major: If true, output will have shape [..., W, N, ...], otherwise
          it will have shape [..., N, W, ...], where x.shape is [..., T, ...].
  
      Returns:
        A tensor containing the stacked slices [0:W, ... T-W:T] from an axis of x.
      """
      sequence_length = x.shape[axis]
      starts = jnp.arange(sequence_length - window + 1)
      ends = jnp.arange(window)
      if time_major:
          idx = (
              starts[None, :] + ends[:, None]
          )  # Output will be [..., W, N, ...].
      else:
          idx = (
              starts[:, None] + ends[None, :]
          )  # Output will be [..., N, W, ...].
      out = jnp.take(x, idx, axis=axis)
      return out
  
file_path: acme/acme/agents/jax/mpo/utils.py
imports: from typing import Callable
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import types
  import distrax
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 3 for calling a global function with example usage: rolling_window(x, window, axis, time_major) and returns: out