code: def topk(scores: mx.nd.NDArray,
           offset: mx.nd.NDArray,
           k: int) -> Tuple[mx.nd.NDArray, mx.nd.NDArray, mx.nd.NDArray]:
      
  
      # Compute the batch size from the offsets and k. We don't know the batch size because it is
      # either 1 (at timestep 1) or k (at timesteps 2+).
      # (batch_size, beam_size * target_vocab_size)
      batch_size = int(offset.shape[-1] / k)
      folded_scores = scores.reshape((batch_size, -1))
  
      # pylint: disable=unbalanced-tuple-unpacking
      values, indices = mx.nd.topk(folded_scores, axis=1, k=k, ret_typ='both', is_ascend=True)
      indices = mx.nd.cast(indices, 'int32').reshape((-1,))
      best_hyp_indices, best_word_indices = mx.nd.unravel_index(indices, shape=(batch_size * k, scores.shape[-1]))
  
      if batch_size > 1:
          # Offsetting the indices to match the shape of the scores matrix
          best_hyp_indices += offset
  
      values = values.reshape((-1, 1))
      return best_hyp_indices, best_word_indices, values
docs: Get the lowest k elements per sentence from a `scores` matrix.
      At the first timestep, the shape of scores is (batch, target_vocabulary_size).
      At subsequent steps, the shape is (batch * k, target_vocabulary_size).
  
      :param scores: Vocabulary scores for the next beam step. (batch_size * beam_size, target_vocabulary_size)
      :param offset: Array (shape: batch_size * k) containing offsets to add to the hypothesis indices in batch decoding.
      :param k: The number of smallest scores to return.
      :return: The row indices, column indices and values of the k smallest items in matrix.