code: def _get_best_word_indices_for_kth_hypotheses(ks: np.ndarray, all_hyp_indices: np.ndarray) -> np.ndarray:
          
          batch_size = ks.shape[0]
          num_steps = all_hyp_indices.shape[1]
          result = np.zeros((batch_size, num_steps - 1), dtype=all_hyp_indices.dtype)
          # first index into the history of the desired hypotheses.
          pointer = all_hyp_indices[ks, -1]
          # for each column/step follow the pointer, starting from the penultimate column/step
          num_steps = all_hyp_indices.shape[1]
          for step in range(num_steps - 2, -1, -1):
              result[:, step] = pointer
              pointer = all_hyp_indices[pointer, step]
          return result
docs: Traverses the matrix of best hypotheses indices collected during beam search in reversed order by
          using the kth hypotheses index as a backpointer.
          Returns an array containing the indices into the best_word_indices collected during beam search to extract
          the kth hypotheses.
  
          :param ks: The kth-best hypotheses to extract. Supports multiple for batch_size > 1. Shape: (batch,).
          :param all_hyp_indices: All best hypotheses indices list collected in beam search. Shape: (batch * beam, steps).
          :return: Array of indices into the best_word_indices collected in beam search
              that extract the kth-best hypothesis. Shape: (batch,).