def as_top_k(
      self,
      k,
      matrix,
      type_name,
      simplify_unitsize_minibatch = True
  ):
    """Return top-k (entity-name, weight) tuples for each row of a matrix.

    Args:
      k: integer
      matrix: numpy ndarray of shape (num_entity_sets, entity_index) with values
        being the weights.
      type_name: string name of the type associated with columns of this matrix.
      simplify_unitsize_minibatch: if true and as_dicts is also true, and the
        minibatch size is 1, then just return the single row dictionary.

    Returns:
      similar to as_dicts but returns lists of top-scoring k items
    """
    if k < 1:
      raise ValueError('k must be positive but it is %d' % k)
    result = []
    num_entity_sets = matrix.shape[0]
    # Find the indices with the highest valued weights.
    top_k_idx = np.flip(np.argsort(matrix, axis=1)[:, -k:], axis=1)
    row_index = np.arange(num_entity_sets).repeat(k)
    column_index = top_k_idx.reshape(-1)
    # Slice, reshape, and sort descending.
    top_k_weights = np.flip(
        np.sort(
            matrix[row_index, column_index].reshape(num_entity_sets, k),
            axis=1),
        axis=1)
    # Convert column indices into entities.
    for indices, weights in zip(top_k_idx, top_k_weights):
      entities = [
          self.get_entity_name(entity_index, type_name)
          for entity_index in indices
      ]
      result.append(list(zip(entities, weights)))
    if simplify_unitsize_minibatch and len(result) == 1:
      return result[0]
    else:
      return result
Return top-k (entity-name, weight) tuples for each row of a matrix.

    Args:
      k: integer
      matrix: numpy ndarray of shape (num_entity_sets, entity_index) with values
        being the weights.
      type_name: string name of the type associated with columns of this matrix.
      simplify_unitsize_minibatch: if true and as_dicts is also true, and the
        minibatch size is 1, then just return the single row dictionary.

    Returns:
      similar to as_dicts but returns lists of top-scoring k items
    
language/nql/nql/__init__.py