code: def chosen_probabs(probab_observations, actions):
    
    B, T = actions.shape  # pylint: disable=invalid-name
    assert (B, T + 1) == probab_observations.shape[:2]
    return probab_observations[np.arange(B)[:, None], np.arange(T), actions]
docs: Picks out the probabilities of the actions along batch and time-steps.
  
    Args:
      probab_observations: ndarray of shape `[B, T+1, A]`, where
        probab_observations[b, t, i] contains the log-probability of action = i at
        the t^th time-step in the b^th trajectory.
      actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which
        action was chosen in the b^th trajectory's t^th time-step.
  
    Returns:
      `[B, T]` ndarray with the log-probabilities of the chosen actions.