code: def nearest_k(similarity_matrix: mx.nd.NDArray,
                query_word_id: int,
                k: int,
                gamma: float = 1.0) -> Iterable[Tuple[int, float]]:
      
      # pylint: disable=unbalanced-tuple-unpacking
      values, indices = mx.nd.topk(mx.nd.softmax(similarity_matrix[query_word_id] / gamma), k=k, ret_typ='both')
      return zip(indices.asnumpy(), values.asnumpy())
docs: Returns values and indices of k items with largest similarity.
  
      :param similarity_matrix: Similarity matrix.
      :param query_word_id: Query word id.
      :param k: Number of closest items to retrieve.
      :param gamma: Parameter to control distribution steepness.
      :return: List of indices and values of k nearest elements.