Python tensorflow.argsort() Examples

def hard_negative_mining(loss, gt_confs, neg_ratio):
    """ Hard negative mining algorithm
        to pick up negative examples for back-propagation
        base on classification loss values
    Args:
        loss: list of classification losses of all default boxes (B, num_default)
        gt_confs: classification targets (B, num_default)
        neg_ratio: negative / positive ratio
    Returns:
        conf_loss: classification loss
        loc_loss: regression loss
    """
    # loss: B x N
    # gt_confs: B x N
    pos_idx = gt_confs > 0
    num_pos = tf.reduce_sum(tf.dtypes.cast(pos_idx, tf.int32), axis=1)
    num_neg = num_pos * neg_ratio

    rank = tf.argsort(loss, axis=1, direction='DESCENDING')
    rank = tf.argsort(rank, axis=1)
    neg_idx = rank < tf.expand_dims(num_neg, 1)

    return pos_idx, neg_idx 

def _to_nd_indices(indices):
    """Returns indices used for tf.gather_nd or tf.scatter_nd.

    Args:
      indices: A `Tensor` of shape [batch_size, size] with integer values. The
        values are the indices of another `Tensor`. For example, `indices` is the
        output of tf.argsort or tf.math.top_k.

    Returns:
      A `Tensor` with shape [batch_size, size, 2] that can be used by tf.gather_nd
      or tf.scatter_nd.

    """
    indices.get_shape().assert_has_rank(2)
    batch_ids = tf.ones_like(indices) * tf.expand_dims(
        tf.range(tf.shape(input=indices)[0]), 1)
    return tf.stack([batch_ids, indices], axis=-1) 

def _to_nd_indices(indices):
  """Returns indices used for tf.gather_nd or tf.scatter_nd.

  Args:
    indices: A `Tensor` of shape [batch_size, size] with integer values. The
      values are the indices of another `Tensor`. For example, `indices` is the
      output of tf.argsort or tf.math.top_k.

  Returns:
    A `Tensor` with shape [batch_size, size, 2] that can be used by tf.gather_nd
    or tf.scatter_nd.

  """
  indices.get_shape().assert_has_rank(2)
  batch_ids = tf.ones_like(indices) * tf.expand_dims(
      tf.range(tf.shape(input=indices)[0]), 1)
  return tf.stack([batch_ids, indices], axis=-1) 

def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
    """
    Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples

    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)
        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        k_expanded = k if isinstance(k, int) else k[:, None]
        exclude_mask = tf.range(vocab_size)[None] >= k_expanded

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

    return {
        'probs': probs,
        'sample': sample,
    } 

def compute_nms(boxes, scores, nms_threshold, limit=200):
    """ Perform Non Maximum Suppression algorithm
        to eliminate boxes with high overlap

    Args:
        boxes: tensor (num_boxes, 4)
               of format (xmin, ymin, xmax, ymax)
        scores: tensor (num_boxes,)
        nms_threshold: NMS threshold
        limit: maximum number of boxes to keep

    Returns:
        idx: indices of kept boxes
    """
    if boxes.shape[0] == 0:
        return tf.constant([], dtype=tf.int32)
    selected = [0]
    idx = tf.argsort(scores, direction='DESCENDING')
    idx = idx[:limit]
    boxes = tf.gather(boxes, idx)

    iou = compute_iou(boxes, boxes)

    while True:
        row = iou[selected[-1]]
        next_indices = row <= nms_threshold
        # iou[:, ~next_indices] = 1.0
        iou = tf.where(
            tf.expand_dims(tf.math.logical_not(next_indices), 0),
            tf.ones_like(iou, dtype=tf.float32),
            iou)

        if not tf.math.reduce_any(next_indices):
            break

        selected.append(tf.argsort(
            tf.dtypes.cast(next_indices, tf.int32), direction='DESCENDING')[0].numpy())

    return tf.gather(idx, selected) 

def build_rt_decoder(self):
        self.build_embeddings()
        start_tokens_sparse = tf.ones(shape=[self.batch_size], dtype=tf.int32) * self.vocab.start_token_id
        with tf.name_scope('beamsearch_decoder'):
            rt_decoder = BeamSearchDecoder(cell=self.cell,
                                           embedding=self.embed_sparse_to_dense,
                                           start_tokens=start_tokens_sparse,
                                           end_token=self.vocab.end_token_id,
                                           initial_state=self.initial_state,
                                           beam_width=self.hparams.beam_width,
                                           output_layer=self.output_layer,
                                           skip_tokens_decoding=self.vocab.skip_tokens,
                                           shrink_vocab=self.hparams.shrink_vocab)
            (hypotheses, input_query_ids, scores) = dynamic_decode(
                decoder=rt_decoder,
                output_time_major=False,
                maximum_iterations=self.hparams.max_length,
                repetition=self.hparams.repetition_penalty)

            sort_ids = tf.argsort(scores, direction="DESCENDING", stable=True, axis=0)
            scores = tf.gather_nd(scores, sort_ids)
            hypotheses = tf.gather_nd(hypotheses, sort_ids)
            input_query_ids = tf.gather_nd(input_query_ids, sort_ids)

            sort_ids = tf.argsort(input_query_ids, direction="ASCENDING", stable=True, axis=0)
            scores = tf.gather_nd(scores, sort_ids)
            hypotheses = tf.gather_nd(hypotheses, sort_ids)
            input_query_ids = tf.gather_nd(input_query_ids, sort_ids)

            speakers = tf.gather_nd(tf.convert_to_tensor(self.features["speaker_ids"]), input_query_ids)
            input_queries = tf.gather_nd(tf.convert_to_tensor(self.features["original_inputs"]), input_query_ids)

            self.rt_hypotheses = tf.identity(hypotheses)
            self.inputs_pred = tf.identity(input_queries)
            self.speakers = tf.identity(speakers)
            self.scores = tf.identity(scores) 

def build_rt_decoder(self):
        self.build_embeddings()
        start_tokens_sparse = tf.ones(shape=[self.batch_size], dtype=tf.int32) * self.vocab.start_token_id
        with tf.name_scope('beamsearch_decoder'):
            rt_decoder = BeamSearchDecoder(cell=self.cell,
                                           embedding=self.embed_sparse_to_dense,
                                           start_tokens=start_tokens_sparse,
                                           end_token=self.vocab.end_token_id,
                                           initial_state=self.initial_state,
                                           beam_width=self.hparams.beam_width,
                                           output_layer=self.output_layer,
                                           skip_tokens_decoding=self.vocab.skip_tokens,
                                           shrink_vocab=self.hparams.shrink_vocab)
            (hypotheses, input_query_ids, scores) = dynamic_decode(
                decoder=rt_decoder,
                output_time_major=False,
                maximum_iterations=self.hparams.max_length,
                repetition=self.hparams.repetition_penalty)

            sort_ids = tf.argsort(
                scores, direction="DESCENDING", stable=True, axis=0)
            scores = tf.gather_nd(scores, sort_ids)
            hypotheses = tf.gather_nd(hypotheses, sort_ids)
            input_query_ids = tf.gather_nd(input_query_ids, sort_ids)

            sort_ids = tf.argsort(
                input_query_ids, direction="ASCENDING", stable=True, axis=0)
            scores = tf.gather_nd(scores, sort_ids)
            hypotheses = tf.gather_nd(hypotheses, sort_ids)
            input_query_ids = tf.gather_nd(input_query_ids, sort_ids)

            input_queries = tf.gather_nd(tf.convert_to_tensor(
                self.features["original_inputs"]), input_query_ids)
            self.rt_hypotheses = tf.identity(hypotheses)
            self.inputs_pred = tf.identity(input_queries)
            self.scores = tf.identity(scores) 

def replace(x, mapping, zero=True):
    """Replace values in tensor `x` using dictionary `mapping`.

    Parameters
    ----------
    x: tensor, values to replace.
    mapping: dict, dictionary mapping original values to new values. Values in
        x equal to a key in the mapping are replaced with the corresponding
        value. Keys and values may overlap.
    zero: boolean, zero values in `x` not in `mapping.keys()`.

    Returns
    -------
    Modified tensor.
    """
    x = tf.cast(x, dtype=tf.int32)
    keys = tf.convert_to_tensor(list(mapping.keys()))
    vals = tf.convert_to_tensor(list(mapping.values()))

    sidx = tf.argsort(keys)
    ks = tf.gather(keys, sidx)
    vs = tf.gather(vals, sidx)

    idx = tf.searchsorted(ks, tf.reshape(x, (-1,)))
    idx = tf.reshape(idx, x.shape)

    # Zero values that are equal to len(vs).
    idx = tf.multiply(idx, tf.cast(tf.not_equal(idx, vs.shape[0]), tf.int32))
    mask = tf.equal(tf.gather(ks, idx), x)
    out = tf.where(mask, tf.gather(vs, idx), x)

    if zero:
        # Zero values in the data array that are not in the mapping values.
        mask = tf.reduce_any(
            tf.equal(tf.expand_dims(out, -1), tf.expand_dims(vals, 0)), -1
        )
        out = tf.multiply(out, tf.cast(mask, tf.int32))

    return out 

def reorder_histogram(bucket_vocab, counts, boundary_size):
  """Return the histogram counts in indexed order, and zero out missing values.

  The count_elements analyzer returns counts in alphanumeric order, only for the
  values that are present. To construct a well-formed histogram, we need to
  rearrange them in numerical order, and fill in the missing values.

  Ex: The data contains values in the following form: [0, 1, 0, 1, 0, 3, 0, 1]
  bucket_indices happen to be the same as these values, and
  count_elements(tf.strings.as_string(bucket_indices)) returns:
    bucket_vocab=['1', '3', '0'],
    counts=[3, 1, 4]

  If boundaries=[0, 1, 2, 3, 4], we expect counts=[4, 3, 0, 1, 0],
  which this function will return.

  Args:
    bucket_vocab: A `Tensor` that names the buckets corresponding to the count
        information returned.
    counts: A `Tensor` that matches the bucket_vocab.
    boundary_size: A scalar that provides information about how big the returned
        counts should be.

  Returns:
    counts: A `Tensor` of size boundary_size corresponding to counts of all
        available buckets.
  """
  if bucket_vocab.dtype == tf.string:
    bucket_vocab = tf.strings.to_number(bucket_vocab, tf.int32)
  # counts/bucket_vocab may be out of order and missing values (empty buckets).
  ordering = tf.argsort(
      tf.concat([bucket_vocab,
                 tf.sets.difference([tf.range(boundary_size)],
                                    [bucket_vocab]).values], axis=-1))
  counts = tf.pad(counts, [[0, boundary_size - tf.size(counts)]])
  return tf.gather(counts, ordering) 

def organize_valid_indices(is_valid, shuffle=True, seed=None):
  """Organizes indices in such a way that valid items appear first.

  Args:
    is_valid: A boolean `Tensor` for entry validity with shape [batch_size,
      list_size].
    shuffle: A boolean indicating whether valid items should be shuffled.
    seed: An int for random seed at the op level. It works together with the
      seed at global graph level together to determine the random number
      generation. See `tf.set_random_seed`.

  Returns:
    A tensor of indices with shape [batch_size, list_size, 2]. The returned
    tensor can be used with `tf.gather_nd` and `tf.scatter_nd` to compose a new
    [batch_size, list_size] tensor. The values in the last dimension are the
    indices for an element in the input tensor.
  """
  with tf.compat.v1.name_scope(name='organize_valid_indices'):
    is_valid = tf.convert_to_tensor(value=is_valid)
    is_valid.get_shape().assert_has_rank(2)
    output_shape = tf.shape(input=is_valid)

    if shuffle:
      values = tf.random.uniform(output_shape, seed=seed)
    else:
      values = (
          tf.ones_like(is_valid, tf.float32) * tf.reverse(
              tf.cast(tf.range(output_shape[1]), dtype=tf.float32), [-1]))

    rand = tf.compat.v1.where(is_valid, values, tf.ones(output_shape) * -1e-6)
    # shape(indices) = [batch_size, list_size]
    indices = tf.argsort(rand, direction='DESCENDING', stable=True)
    return _to_nd_indices(indices) 

def sorted_ranks(scores, shuffle_ties=True, seed=None):
  """Returns an int `Tensor` as the ranks (1-based) after sorting scores.

  Example: Given scores = [[1.0, 3.5, 2.1]], the returned ranks will be [[3, 1,
  2]]. It means that scores 1.0 will be ranked at position 3, 3.5 will be ranked
  at position 1, and 2.1 will be ranked at position 2.

  Args:
    scores: A `Tensor` of shape [batch_size, list_size] representing the
      per-example scores.
    shuffle_ties: See `sort_by_scores`.
    seed: See `sort_by_scores`.

  Returns:
    A 1-based int `Tensor`s as the ranks.
  """
  with tf.compat.v1.name_scope(name='sorted_ranks'):
    batch_size, list_size = tf.unstack(tf.shape(input=scores))
    # The current position in the list for each score.
    positions = tf.tile(tf.expand_dims(tf.range(list_size), 0), [batch_size, 1])
    # For score [[1.0, 3.5, 2.1]], sorted_positions are [[1, 2, 0]], meaning the
    # largest score is at position 1, the 2nd is at position 2 and 3rd is at
    # position 0.
    sorted_positions = sort_by_scores(
        scores, [positions], shuffle_ties=shuffle_ties, seed=seed)[0]
    # The indices of sorting sorted_positions will be [[2, 0, 1]] and ranks are
    # 1-based and thus are [[3, 1, 2]].
    ranks = tf.argsort(sorted_positions) + 1
    return ranks 

def sort_by_scores(scores,
                   features_list,
                   topn=None,
                   shuffle_ties=True,
                   seed=None):
  """Sorts example features according to per-example scores.

  Args:
    scores: A `Tensor` of shape [batch_size, list_size] representing the
      per-example scores.
    features_list: A list of `Tensor`s with the same shape as scores to be
      sorted.
    topn: An integer as the cutoff of examples in the sorted list.
    shuffle_ties: A boolean. If True, randomly shuffle before the sorting.
    seed: The ops-level random seed used when `shuffle_ties` is True.

  Returns:
    A list of `Tensor`s as the list of sorted features by `scores`.
  """
  with tf.compat.v1.name_scope(name='sort_by_scores'):
    scores = tf.cast(scores, tf.float32)
    scores.get_shape().assert_has_rank(2)
    list_size = tf.shape(input=scores)[1]
    if topn is None:
      topn = list_size
    topn = tf.minimum(topn, list_size)
    shuffle_ind = None
    if shuffle_ties:
      shuffle_ind = _to_nd_indices(
          tf.argsort(
              tf.random.uniform(tf.shape(input=scores), seed=seed),
              stable=True))
      scores = tf.gather_nd(scores, shuffle_ind)
    _, indices = tf.math.top_k(scores, topn, sorted=True)
    nd_indices = _to_nd_indices(indices)
    if shuffle_ind is not None:
      nd_indices = tf.gather_nd(shuffle_ind, nd_indices)
    return [tf.gather_nd(f, nd_indices) for f in features_list] 

def organize_valid_indices(is_valid, shuffle=True, seed=None):
    """Organizes indices in such a way that valid items appear first.

    Args:
      is_valid: A boolen `Tensor` for entry validity with shape [batch_size,
        list_size].
      shuffle: A boolean indicating whether valid items should be shuffled.
      seed: An int for random seed at the op level. It works together with the
        seed at global graph level together to determine the random number
        generation. See `tf.set_random_seed`.

    Returns:
      A tensor of indices with shape [batch_size, list_size, 2]. The returned
      tensor can be used with `tf.gather_nd` and `tf.scatter_nd` to compose a new
      [batch_size, list_size] tensor. The values in the last dimension are the
      indices for an element in the input tensor.
    """
    with tf.compat.v1.name_scope(name='organize_valid_indices'):
        is_valid = tf.convert_to_tensor(value=is_valid)
        is_valid.get_shape().assert_has_rank(2)
        output_shape = tf.shape(input=is_valid)

        if shuffle:
            values = tf.random.uniform(output_shape, seed=seed)
        else:
            values = (
                tf.ones_like(is_valid, tf.float32) * tf.reverse(
                    tf.cast(tf.range(output_shape[1]), dtype=tf.float32), [-1]))

        rand = tf.compat.v1.where(
            is_valid, values, tf.ones(output_shape) * -1e-6)
        # shape(indices) = [batch_size, list_size]
        indices = tf.argsort(rand, direction='DESCENDING', stable=True)
        return _to_nd_indices(indices) 

def sorted_ranks(scores, shuffle_ties=True, seed=None):
    """Returns an int `Tensor` as the ranks (1-based) after sorting scores.

    Example: Given scores = [[1.0, 3.5, 2.1]], the returned ranks will be [[3, 1,
    2]]. It means that scores 1.0 will be ranked at position 3, 3.5 will be ranked
    at position 1, and 2.1 will be ranked at position 2.

    Args:
      scores: A `Tensor` of shape [batch_size, list_size] representing the
        per-example scores.
      shuffle_ties: See `sort_by_scores`.
      seed: See `sort_by_scores`.

    Returns:
      A 1-based int `Tensor`s as the ranks.
    """
    with tf.compat.v1.name_scope(name='sorted_ranks'):
        batch_size, list_size = tf.unstack(tf.shape(input=scores))
        # The current position in the list for each score.
        positions = tf.tile(
            tf.expand_dims(
                tf.range(list_size), 0), [
                batch_size, 1])
        # For score [[1.0, 3.5, 2.1]], sorted_positions are [[1, 2, 0]], meaning the
        # largest score is at poistion 1, the second is at postion 2 and third is at
        # position 0.
        sorted_positions = sort_by_scores(
            scores, [positions], shuffle_ties=shuffle_ties, seed=seed)[0]
        # The indices of sorting sorted_postions will be [[2, 0, 1]] and ranks are
        # 1-based and thus are [[3, 1, 2]].
        ranks = tf.argsort(sorted_positions) + 1
        return ranks 

def sort_by_scores(scores,
                   features_list,
                   topn=None,
                   shuffle_ties=True,
                   seed=None):
    """Sorts example features according to per-example scores.

    Args:
      scores: A `Tensor` of shape [batch_size, list_size] representing the
        per-example scores.
      features_list: A list of `Tensor`s with the same shape as scores to be
        sorted.
      topn: An integer as the cutoff of examples in the sorted list.
      shuffle_ties: A boolean. If True, randomly shuffle before the sorting.
      seed: The ops-level random seed used when `shuffle_ties` is True.

    Returns:
      A list of `Tensor`s as the list of sorted features by `scores`.
    """
    with tf.compat.v1.name_scope(name='sort_by_scores'):
        scores = tf.cast(scores, tf.float32)
        scores.get_shape().assert_has_rank(2)
        list_size = tf.shape(input=scores)[1]
        if topn is None:
            topn = list_size
        topn = tf.minimum(topn, list_size)
        shuffle_ind = None
        if shuffle_ties:
            shuffle_ind = _to_nd_indices(
                tf.argsort(
                    tf.random.uniform(tf.shape(input=scores), seed=seed),
                    stable=True))
            scores = tf.gather_nd(scores, shuffle_ind)
        _, indices = tf.math.top_k(scores, topn, sorted=True)
        nd_indices = _to_nd_indices(indices)
        if shuffle_ind is not None:
            nd_indices = tf.gather_nd(shuffle_ind, nd_indices)
        return [tf.gather_nd(f, nd_indices) for f in features_list] 

def get_top_elements(list_of_elements, max_user_contribution):
  """Gets the top max_user_contribution words from the input list.

  Note that the returned set of top words will not necessarily be sorted.

  Args:
    list_of_elements: A tensor containing a list of elements.
    max_user_contribution: The maximum number of elements to keep.

  Returns:
    A tensor of a list of strings.
    If the total number of unique words is less than or equal to
    max_user_contribution, returns the set of unique words.
  """
  words, _, counts = tf.unique_with_counts(list_of_elements)
  if tf.size(words) > max_user_contribution:
    # This logic is influenced by the focus on global heavy hitters and
    # thus implements clipping by chopping the tail of the distribution
    # of the words as present on a single client. Another option could
    # be to provide pick max_words_per_user random words out of the unique
    # words present locally.
    top_indices = tf.argsort(
        counts, axis=-1, direction='DESCENDING')[:max_user_contribution]
    top_words = tf.gather(words, top_indices)
    return top_words
  return words 

def _apply(self, words):
    if self.max_distance == 0:
      return tf.identity(words)
    num_words = tf.shape(words)[0]
    offset = tf.random.uniform([num_words], maxval=1) * (self.max_distance + 1)
    offset = tf.cast(offset, num_words.dtype)
    new_pos = tf.argsort(tf.range(num_words) + offset)
    return tf.gather(words, new_pos) 

def sort_key_val(t1, t2, dim=-1):
    values = tf.sort(t1, axis=dim)
    t2 = tf.broadcast_to(t2, t1.shape)
    return values, tf.gather(t2, tf.argsort(t1, axis=dim), axis=dim) 

def nucleus_sampling(logits, vocab_size, p=0.9, 
					input_ids=None, input_ori_ids=None,
					**kargs):
	input_shape_list = bert_utils.get_shape_list(logits, expected_rank=[2,3])
	if len(input_shape_list) == 3:
		logits = tf.reshape(logits, (-1, vocab_size))
	probs = tf.nn.softmax(logits, axis=-1)
	# [batch_size, seq, vocab_perm]
	# indices = tf.argsort(probs, direction='DESCENDING')
	indices = tf.contrib.framework.argsort(probs, direction='DESCENDING')

	cumulative_probabilities = tf.math.cumsum(tf.batch_gather(probs, indices), axis=-1, exclusive=False)
	
	# find the top pth index to cut off. careful we don't want to cutoff everything!
	# result will be [batch_size, seq, vocab_perm]
	exclude_mask = tf.logical_not(
	tf.logical_or(cumulative_probabilities < p, tf.range(vocab_size)[None] < 1))
	exclude_mask = tf.cast(exclude_mask, tf.float32)

	indices_v1 = tf.contrib.framework.argsort(indices)
	exclude_mask = reorder(exclude_mask, tf.cast(indices_v1, dtype=tf.int32))
	if len(input_shape_list) == 3:
		exclude_mask = tf.reshape(exclude_mask, input_shape_list)
		# logits = tf.reshape(logits, input_shape_list)

	if input_ids is not None and input_ori_ids is not None:
		exclude_mask, input_ori_ids = get_extra_mask(
								input_ids, input_ori_ids, 
								exclude_mask, vocab_size,
								**kargs)

		return [exclude_mask, input_ori_ids]
	else:
		return [exclude_mask] 

def gumbel_softmax(logits, temperature, gumbel_samples=None, samples=1, greedy=False): 
	""" Draw a sample from the Gumbel-Softmax distribution"""
	input_shape_list = bert_utils.get_shape_list(logits, expected_rank=2)
	if samples > 1:
		logits = tf.expand_dims(logits, -1)
	if gumbel_samples is None:
		gumbel_samples = sample_gumbel(input_shape_list, samples)
	if greedy:
		tf.logging.info("==apply greedy based sampling and discrete relax==")
		# if int(tf.__version__.split(".")[1]) < 15:
		# 	if not use_tpu:
		# 		logits_index = tf.contrib.framework.argsort(logits, axis=1)
		# 		gumbel_samples_sorted = tf.contrib.framework.sort(gumbel_samples, axis=1)
		# 		gumbel_samples_sorted = reorder(gumbel_samples_sorted, logits_index)
		# 	else:

		# else:
		# 	logits_index = tf.argsort(logits, axis=1)
		# 	gumbel_samples_sorted = tf.sort(gumbel_samples, axis=1)
		# 	gumbel_samples_sorted = reorder(gumbel_samples_sorted, logits_index)

		gumbel_samples = reorder_approximate(logits, gumbel_samples)
		y = logits + gumbel_samples
		return [tf.exp(tf.nn.log_softmax(y / temperature, axis=1)),
				y]
	else:
		y = logits + gumbel_samples
		tf.logging.info("==apply sampling based sampling and discrete relax==")
		return [tf.exp(tf.nn.log_softmax(y / temperature, axis=1)), 
				y] 

def call(self, inputs):
        if self.data_mode == 'disjoint':
            X, I = inputs
            X = ops.disjoint_signal_to_batch(X, I)
        else:
            X = inputs
            if self.data_mode == 'single':
                X = tf.expand_dims(X, 0)

        N = tf.shape(X)[-2]
        sort_perm = tf.argsort(X[..., -1], direction='DESCENDING')
        X_sorted = tf.gather(X, sort_perm, axis=-2, batch_dims=1)

        def truncate():
            _X_out = X_sorted[..., : self.k, :]
            return _X_out

        def pad():
            padding = [[0, 0], [0, self.k - N], [0, 0]]
            _X_out = tf.pad(X_sorted, padding)
            return _X_out

        X_out = tf.cond(tf.less_equal(self.k, N), truncate, pad)

        if self.data_mode == 'single':
            X_out = tf.squeeze(X_out, [0])
            X_out.set_shape((self.k, self.F))
        elif self.data_mode == 'batch' or self.data_mode == 'disjoint':
            X_out.set_shape((None, self.k, self.F))

        return X_out 

def compute_top_k_and_ndcg(logits,              # type: tf.Tensor
                           duplicate_mask,      # type: tf.Tensor
                           match_mlperf=False   # type: bool
                          ):
  """Compute inputs of metric calculation.

  Args:
    logits: A tensor containing the predicted logits for each user. The shape
      of logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits
      for a user are grouped, and the first element of the group is the true
      element.
    duplicate_mask: A vector with the same shape as logits, with a value of 1
      if the item corresponding to the logit at that position has already
      appeared for that user.
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
    is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
    logits_by_user which has size
    (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
  """
  logits_by_user = tf.reshape(logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1))
  duplicate_mask_by_user = tf.cast(
      tf.reshape(duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)),
      logits_by_user.dtype)

  if match_mlperf:
    # Set duplicate logits to the min value for that dtype. The MLPerf
    # reference dedupes during evaluation.
    logits_by_user *= (1 - duplicate_mask_by_user)
    logits_by_user += duplicate_mask_by_user * logits_by_user.dtype.min

  # Determine the location of the first element in each row after the elements
  # are sorted.
  sort_indices = tf.argsort(
      logits_by_user, axis=1, direction="DESCENDING")

  # Use matrix multiplication to extract the position of the true item from the
  # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
  # perform matrix multiplications very quickly. This is similar to np.argwhere.
  # However this is a special case because the target will only appear in
  # sort_indices once.
  one_hot_position = tf.cast(tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES),
                             tf.int32)
  sparse_positions = tf.multiply(
      one_hot_position, tf.range(logits_by_user.shape[1])[tf.newaxis, :])
  position_vector = tf.reduce_sum(sparse_positions, axis=1)

  in_top_k = tf.cast(tf.less(position_vector, rconst.TOP_K), tf.float32)
  ndcg = tf.math.log(2.) / tf.math.log(
      tf.cast(position_vector, tf.float32) + 2)
  ndcg *= in_top_k

  # If a row is a padded row, all but the first element will be a duplicate.
  metric_weights = tf.not_equal(tf.reduce_sum(duplicate_mask_by_user, axis=1),
                                rconst.NUM_EVAL_NEGATIVES)

  return in_top_k, ndcg, metric_weights, logits_by_user 

def de_noise(counts, noise, ratio=0.9):
  """Returns a float `Tensor` as the de-noised `counts`.

  The implementation is based on the the paper by Zhang and Xu: "Fast Exact
  Maximum Likelihood Estimation for Mixture of Language Models." It assumes that
  the observed `counts` are generated from a mixture of `noise` and the true
  distribution: `ratio * noise_distribution + (1 - ratio) * true_distribution`,
  where the contribution of `noise` is controlled by `ratio`. This method
  returns the true distribution.

  Args:
    counts: A 2-D `Tensor` representing the observations. All values should be
      nonnegative.
    noise: A 2-D `Tensor` representing the noise distribution. This should be
      the same shape as `counts`. All values should be positive and are
      normalized to a simplex per row.
    ratio: A float in (0, 1) representing the contribution from noise.

  Returns:
    A 2-D float `Tensor` and each row is a simplex.
  Raises:
    ValueError: if `ratio` is not in (0,1).
    InvalidArgumentError: if any of `counts` is negative or any of `noise` is
    not positive.
  """
  if not 0 < ratio < 1:
    raise ValueError('ratio should be in (0, 1), but get {}'.format(ratio))
  odds = (1 - ratio) / ratio

  counts = tf.cast(counts, dtype=tf.float32)
  noise = tf.cast(noise, dtype=tf.float32)

  counts.get_shape().assert_has_rank(2)
  noise.get_shape().assert_has_rank(2)
  noise.get_shape().assert_is_compatible_with(counts.get_shape())

  with tf.compat.v1.name_scope(name='de_noise'):
    counts_nonneg = tf.compat.v1.assert_greater_equal(counts, 0.)
    noise_pos = tf.compat.v1.assert_greater(noise, 0.)
    with tf.control_dependencies([counts_nonneg, noise_pos]):
      # Normalize noise to be a simplex per row.
      noise = noise / tf.reduce_sum(noise, axis=1, keepdims=True)
      sorted_idx = tf.argsort(
          counts / noise, direction='DESCENDING', stable=True)
      nd_indices = _to_nd_indices(sorted_idx)
      sorted_counts = tf.gather_nd(counts, nd_indices)
      sorted_noise = tf.gather_nd(noise, nd_indices)
      # Decide whether an entry will have a positive value or 0.
      is_pos = tf.cast(
          (odds + tf.cumsum(sorted_noise, axis=1)) /
          tf.cumsum(sorted_counts, axis=1) > sorted_noise / sorted_counts,
          tf.float32)
      # The lambda in the paper above, which is the lagrangian multiplier for
      # the simplex constraint on the variables.
      lagrangian_multiplier = tf.reduce_sum(
          sorted_counts * is_pos, axis=1, keepdims=True) / (1 + tf.reduce_sum(
              sorted_noise * is_pos, axis=1, keepdims=True) / odds)
      res = (sorted_counts / lagrangian_multiplier -
             sorted_noise / odds) * is_pos
      return tf.scatter_nd(nd_indices, res, shape=tf.shape(counts)) 

def compute_top_k_and_ndcg(logits,              # type: tf.Tensor
                           duplicate_mask,      # type: tf.Tensor
                           match_mlperf=False   # type: bool
                          ):
  """Compute inputs of metric calculation.

  Args:
    logits: A tensor containing the predicted logits for each user. The shape
      of logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits
      for a user are grouped, and the first element of the group is the true
      element.
    duplicate_mask: A vector with the same shape as logits, with a value of 1
      if the item corresponding to the logit at that position has already
      appeared for that user.
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
    is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
    logits_by_user which has size
    (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
  """
  logits_by_user = tf.reshape(logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1))
  duplicate_mask_by_user = tf.cast(
      tf.reshape(duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)),
      tf.float32)

  if match_mlperf:
    # Set duplicate logits to the min value for that dtype. The MLPerf
    # reference dedupes during evaluation.
    logits_by_user *= (1 - duplicate_mask_by_user)
    logits_by_user += duplicate_mask_by_user * logits_by_user.dtype.min

  # Determine the location of the first element in each row after the elements
  # are sorted.
  sort_indices = tf.argsort(
      logits_by_user, axis=1, direction="DESCENDING")

  # Use matrix multiplication to extract the position of the true item from the
  # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
  # perform matrix multiplications very quickly. This is similar to np.argwhere.
  # However this is a special case because the target will only appear in
  # sort_indices once.
  one_hot_position = tf.cast(tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES),
                             tf.int32)
  sparse_positions = tf.multiply(
      one_hot_position, tf.range(logits_by_user.shape[1])[tf.newaxis, :])
  position_vector = tf.reduce_sum(sparse_positions, axis=1)

  in_top_k = tf.cast(tf.less(position_vector, rconst.TOP_K), tf.float32)
  ndcg = tf.math.log(2.) / tf.math.log(
      tf.cast(position_vector, tf.float32) + 2)
  ndcg *= in_top_k

  # If a row is a padded row, all but the first element will be a duplicate.
  metric_weights = tf.not_equal(tf.reduce_sum(duplicate_mask_by_user, axis=1),
                                rconst.NUM_EVAL_NEGATIVES)

  return in_top_k, ndcg, metric_weights, logits_by_user 

def segment_top_k(x, I, ratio, top_k_var):
    """
    Returns indices to get the top K values in x segment-wise, according to
    the segments defined in I. K is not fixed, but it is defined as a ratio of
    the number of elements in each segment.
    :param x: a rank 1 Tensor;
    :param I: a rank 1 Tensor with segment IDs for x;
    :param ratio: float, ratio of elements to keep for each segment;
    :param top_k_var: a tf.Variable created without shape validation (i.e.,
    `tf.Variable(0.0, validate_shape=False)`);
    :return: a rank 1 Tensor containing the indices to get the top K values of
    each segment in x.
    """
    I = tf.cast(I, tf.int32)
    num_nodes = tf.math.segment_sum(tf.ones_like(I), I)  # Number of nodes in each graph
    cumsum = tf.cumsum(num_nodes)  # Cumulative number of nodes (A, A+B, A+B+C)
    cumsum_start = cumsum - num_nodes  # Start index of each graph
    n_graphs = tf.shape(num_nodes)[0]  # Number of graphs in batch
    max_n_nodes = tf.reduce_max(num_nodes)  # Order of biggest graph in batch
    batch_n_nodes = tf.shape(I)[0]  # Number of overall nodes in batch
    to_keep = tf.math.ceil(ratio * tf.cast(num_nodes, tf.float32))
    to_keep = tf.cast(to_keep, I.dtype)  # Nodes to keep in each graph

    index = tf.range(batch_n_nodes)
    index = (index - tf.gather(cumsum_start, I)) + (I * max_n_nodes)

    y_min = tf.reduce_min(x)
    dense_y = tf.ones((n_graphs * max_n_nodes,))
    # subtract 1 to ensure that filler values do not get picked
    dense_y = dense_y * tf.cast(y_min - 1, dense_y.dtype)
    dense_y = tf.cast(dense_y, top_k_var.dtype)
    # top_k_var is a variable with unknown shape defined in the elsewhere
    top_k_var.assign(dense_y)
    dense_y = tf.tensor_scatter_nd_update(top_k_var, index[..., None], tf.cast(x, top_k_var.dtype))
    dense_y = tf.reshape(dense_y, (n_graphs, max_n_nodes))

    perm = tf.argsort(dense_y, direction='DESCENDING')
    perm = perm + cumsum_start[:, None]
    perm = tf.reshape(perm, (-1,))

    to_rep = tf.tile(tf.constant([1., 0.]), (n_graphs,))
    rep_times = tf.reshape(tf.concat((to_keep[:, None], (max_n_nodes - to_keep)[:, None]), -1), (-1,))
    mask = repeat(to_rep, rep_times)

    perm = tf.boolean_mask(perm, mask)

    return perm 

def _top_p_sample(logits, ignore_ids=None, num_samples=1, p=0.9):
    """
    Does top-p sampling. if ignore_ids is on, then we will zero out those logits.
    :param logits: [batch_size, vocab_size] tensor
    :param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
                        like padding maybe
    :param p: topp threshold to use, either a float or a [batch_size] vector
    :return: [batch_size, num_samples] samples

    # TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
    """
    with tf.variable_scope('top_p_sample'):
        batch_size, vocab_size = get_shape_list(logits, expected_rank=2)

        probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                              axis=-1)

        if isinstance(p, float) and p > 0.999999:
            # Don't do top-p sampling in this case
            print("Top-p sampling DISABLED", flush=True)
            return {
                'probs': probs,
                'sample': tf.random.categorical(
                    logits=logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
                    num_samples=num_samples, dtype=tf.int32),
            }

        # [batch_size, vocab_perm]
        indices = tf.argsort(probs, direction='DESCENDING')
        cumulative_probabilities = tf.math.cumsum(tf.batch_gather(probs, indices), axis=-1, exclusive=False)

        # find the top pth index to cut off. careful we don't want to cutoff everything!
        # result will be [batch_size, vocab_perm]
        p_expanded = p if isinstance(p, float) else p[:, None]
        exclude_mask = tf.logical_not(
            tf.logical_or(cumulative_probabilities < p_expanded, tf.range(vocab_size)[None] < 1))

        # OPTION A - sample in the sorted space, then unsort.
        logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
        sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
        sample = tf.batch_gather(indices, sample_perm)

        # OPTION B - unsort first - Indices need to go back to 0 -> N-1 -- then sample
        # unperm_indices = tf.argsort(indices, direction='ASCENDING')
        # include_mask_unperm = tf.batch_gather(include_mask, unperm_indices)
        # logits_to_use = logits - (1 - tf.cast(include_mask_unperm, tf.float32)) * 1e10
        # sample = tf.random.categorical(logits=logits_to_use, num_samples=num_samples, dtype=tf.int32)

    return {
        'probs': probs,
        'sample': sample,
    } 

def compute_top_k_and_ndcg(logits: tf.Tensor,
                           duplicate_mask: tf.Tensor,
                           match_mlperf: bool = False):
  """Compute inputs of metric calculation.

  Args:
    logits: A tensor containing the predicted logits for each user. The shape of
      logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits for a
      user are grouped, and the first element of the group is the true element.
    duplicate_mask: A vector with the same shape as logits, with a value of 1 if
      the item corresponding to the logit at that position has already appeared
      for that user.
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
    is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
    logits_by_user which has size
    (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
  """
  logits_by_user = tf.reshape(logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1))
  duplicate_mask_by_user = tf.cast(
      tf.reshape(duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)),
      logits_by_user.dtype)

  if match_mlperf:
    # Set duplicate logits to the min value for that dtype. The MLPerf
    # reference dedupes during evaluation.
    logits_by_user *= (1 - duplicate_mask_by_user)
    logits_by_user += duplicate_mask_by_user * logits_by_user.dtype.min

  # Determine the location of the first element in each row after the elements
  # are sorted.
  sort_indices = tf.argsort(
      logits_by_user, axis=1, direction="DESCENDING")

  # Use matrix multiplication to extract the position of the true item from the
  # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
  # perform matrix multiplications very quickly. This is similar to np.argwhere.
  # However this is a special case because the target will only appear in
  # sort_indices once.
  one_hot_position = tf.cast(tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES),
                             tf.int32)
  sparse_positions = tf.multiply(
      one_hot_position, tf.range(logits_by_user.shape[1])[tf.newaxis, :])
  position_vector = tf.reduce_sum(sparse_positions, axis=1)

  in_top_k = tf.cast(tf.less(position_vector, rconst.TOP_K), tf.float32)
  ndcg = tf.math.log(2.) / tf.math.log(
      tf.cast(position_vector, tf.float32) + 2)
  ndcg *= in_top_k

  # If a row is a padded row, all but the first element will be a duplicate.
  metric_weights = tf.not_equal(tf.reduce_sum(duplicate_mask_by_user, axis=1),
                                rconst.NUM_EVAL_NEGATIVES)

  return in_top_k, ndcg, metric_weights, logits_by_user