code: def extract_subsequences(
      trajectories: chex.Array,
      start_indices: chex.Array,
      subsequence_len: int = 1,
      max_valid_start_idx: Optional[int] = None,
  ) -> chex.Array:
      """Extract (potentially overlapping) subsequences from batch of trajectories.
  
      WARNING: If `max_valid_start_idx` is not provided, or incorrectly set,
      the function cannot check the validity of the chosen `start_idx` and no error
      will be raised if indexing outside the data boundaries.
  
      Args:
        trajectories: A batch of trajectories, shape `[T, B, ...]`.
        start_indices: Time indices of start points, shape `[B, num_start_indices]`.
        subsequence_len: The length of subsequences extracted from `trajectories`.
        max_valid_start_idx: the maximum valid start index, therefore the
          `start_indices` should be from {0, ..., max_valid_start_idx}.
  
      Returns:
        A batch of subsequences, with
        `trajectories[start_indices[i, j]:start_indices[i, j] + n]` for each start
        index. Output shape is: `[subsequence_len, B, num_start_indices, ...]`.
      """
      if max_valid_start_idx is not None:
          min_len = max_valid_start_idx + subsequence_len
          traj_len = trajectories.shape[0]
          if traj_len < min_len:
              raise AssertionError(
                  f"Expected len >= {min_len}, but trajectories length is: {traj_len}."
              )
  
      batch_size = start_indices.shape[0]
      batch_range = jnp.arange(batch_size)
      num_subs = start_indices.shape[1]
      idx_arr = (
          jnp.arange(subsequence_len)[:, None, None]
          * jnp.ones(
              (subsequence_len, batch_size, num_subs), dtype=jnp.int32
          )
          + start_indices
      )
      return trajectories[idx_arr, batch_range[None, :, None], ...]
  
imports: from typing import Optional
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: extract_subsequences(trajectories, start_indices, subsequence_len, max_valid_start_idx) and returns: trajectories