code: def sample_start_indices(
      rng_key: chex.PRNGKey,
      batch_size: int,
      num_start_indices: int,
      max_valid_start_idx: int,
  ) -> chex.Array:
      """Sampling `batch_size x num_start_indices` starting indices.
  
      Args:
        rng_key: a pseudo random number generator's key.
        batch_size: the size of the batch of trajectories to index in.
        num_start_indices: how many starting points per trajectory in the batch.
        max_valid_start_idx: maximum valid time index for all starting points.
  
      Returns:
        an array of starting points with shape `[B, num_start_indices]`
      """
  
      @jax.vmap
      def _vchoose(key, entries):
          return jax.random.choice(
              key, entries, shape=(num_start_indices,), replace=False
          )
  
      rollout_window = jnp.arange(max_valid_start_idx + 1)
      return _vchoose(
          jax.random.split(rng_key, batch_size),
          jnp.tile(rollout_window, (batch_size, 1)),
      )
  
imports: from typing import Optional
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: sample_start_indices(rng_key, batch_size, num_start_indices, max_valid_start_idx) and returns: _vchoose, jax