Python tensorflow.sort() Examples

The following are 15 code examples of tensorflow.sort(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module tensorflow , or try the search function .
Example #1
Source File: core.py    From lm-human-preferences with MIT License 6 votes vote down vote up
def take_top_p_logits(logits, p):
    """Nucleus sampling"""
    batch, sequence, _ = logits.shape.as_list()
    sorted_logits = tf.sort(logits, direction='DESCENDING', axis=-1)
    cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)
    indices = tf.stack([
        tf.range(0, batch)[:, tf.newaxis],
        tf.range(0, sequence)[tf.newaxis, :],
        # number of indices to include
        tf.maximum(tf.reduce_sum(tf.cast(cumulative_probs <= p, tf.int32), axis=-1) - 1, 0),
    ], axis=-1)
    min_values = tf.gather_nd(sorted_logits, indices)
    return tf.where(
        logits < min_values,
        tf.ones_like(logits) * -1e10,
        logits,
    ) 
Example #2
Source File: glob.py    From dgl with Apache License 2.0 6 votes vote down vote up
def call(self, graph, feat):
        r"""Compute sort pooling.

        Parameters
        ----------
        graph : DGLGraph
            The graph.
        feat : tf.Tensor
            The input feature with shape :math:`(N, D)` where
            :math:`N` is the number of nodes in the graph.

        Returns
        -------
        tf.Tensor
            The output feature with shape :math:`(B, k * D)`, where
            :math:`B` refers to the batch size.
        """
        with graph.local_scope():
            # Sort the feature of each node in ascending order.
            feat = tf.sort(feat, -1)
            graph.ndata['h'] = feat
            # Sort nodes according to their last features.
            ret = tf.reshape(topk_nodes(graph, 'h', self.k, idx=-1)[0], (
                -1, self.k * feat.shape[-1]))
            return ret 
Example #3
Source File: tf_utils.py    From transform with Apache License 2.0 6 votes vote down vote up
def _iteration_l_moments_dense(
    current_index, l1_sum, l2_sum, l3_sum, l4_sum, l1_factors, l2_factors,
    l3_factors, l4_factors, x_rank_2):
  """Process one column of a `Tensor` and updates L-moments variables."""
  current_x = x_rank_2[:, current_index]
  sorted_x = tf.sort(current_x)

  dim_1 = tf.shape(x_rank_2)[1]
  new_l1_sum = l1_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l1_factors), axis=0)], [dim_1])
  new_l2_sum = l2_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l2_factors), axis=0)], [dim_1])
  new_l3_sum = l3_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l3_factors), axis=0)], [dim_1])
  new_l4_sum = l4_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l4_factors), axis=0)], [dim_1])
  return (tf.add(current_index, 1),
          new_l1_sum, new_l2_sum, new_l3_sum, new_l4_sum, l1_factors,
          l2_factors, l3_factors, l4_factors, x_rank_2) 
Example #4
Source File: sample.py    From storybro with MIT License 6 votes vote down vote up
def top_p_logits(logits, p):
    """Nucleus sampling"""
    batch, _ = logits.shape.as_list()
    sorted_logits = tf.sort(logits, direction="DESCENDING", axis=-1)
    cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)
    indices = tf.stack(
        [
            tf.range(0, batch),
            # number of indices to include
            tf.maximum(
                tf.reduce_sum(tf.cast(cumulative_probs <= p, tf.int32), axis=-1) - 1, 0
            ),
        ],
        axis=-1,
    )
    min_values = tf.gather_nd(sorted_logits, indices)
    return tf.where(logits < min_values, tf.ones_like(logits) * -1e10, logits,) 
Example #5
Source File: data_utils.py    From models with Apache License 2.0 6 votes vote down vote up
def _single_token_mask(inputs, tgt_len, num_predict):
  """Sample individual tokens as prediction targets."""
  all_indices = tf.range(tgt_len, dtype=tf.int64)
  non_func_mask = tf.logical_and(
      tf.not_equal(inputs, SEP_ID),
      tf.not_equal(inputs, CLS_ID))
  non_func_indices = tf.boolean_mask(all_indices, non_func_mask)

  masked_pos = tf.random.shuffle(non_func_indices)
  masked_pos = tf.sort(masked_pos[:num_predict])
  target_mask = tf.sparse_to_dense(
      sparse_indices=masked_pos,
      output_shape=[tgt_len],
      sparse_values=1.0,
      default_value=0.0)

  is_masked = tf.cast(target_mask, tf.bool)

  return is_masked, target_mask 
Example #6
Source File: tensor.py    From dgl with Apache License 2.0 5 votes vote down vote up
def sort_1d(input):
    return tf.sort(input), tf.cast(tf.argsort(input), dtype=tf.int64) 
Example #7
Source File: token_generator_gumbel.py    From BERT with Apache License 2.0 5 votes vote down vote up
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] 
Example #8
Source File: reformer_utils.py    From BERT with Apache License 2.0 5 votes vote down vote up
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) 
Example #9
Source File: noise.py    From OpenNMT-tf with MIT License 5 votes vote down vote up
def _apply(self, words):
    if self.count == 0:
      return tf.identity(words)
    num_words = tf.shape(words)[0]
    indices = tf.range(num_words)
    shuffle_indices = tf.random.shuffle(indices)
    keep_count = tf.maximum(num_words - self.count, 1)
    keep_indices = tf.sort(shuffle_indices[:keep_count])
    return tf.gather(words, keep_indices) 
Example #10
Source File: realmix.py    From realmix with Apache License 2.0 5 votes vote down vote up
def percent_confidence_mask_unsup(self, logits_y, labels_y, loss_l2u):
        # Adapted from google-research/uda/image/main.py

        # This function masks the unsupervised predictions that are below
        # a set confidence threshold. # Note the following will only work
        # using MSE loss and not KL-divergence.

        # Calculate largest predicted probability for each image.
        unsup_prob = tf.nn.softmax(logits_y, axis=-1)
        largest_prob = tf.reduce_max(unsup_prob, axis=-1)

        # Get the indices of the bottom x% of probabilities and mask those out.
        # In other words, get the probability of the image with the x%*#numofsamples
        # lowest probability and use that as the mask.

        # Calculate the current confidence_mask value using the specified schedule:
        sorted_probs = tf.sort(largest_prob, axis=-1, direction='ASCENDING')
        sort_index = tf.math.multiply(tf.to_float(tf.shape(sorted_probs)[0]), FLAGS.percent_mask)
        curr_confidence_mask = tf.slice(sorted_probs, [tf.to_int64(sort_index)], [1])

        # Mask the loss for images that don't contain a predicted
        # probability above the threshold.
        loss_mask = tf.cast(tf.greater(largest_prob, curr_confidence_mask), tf.float32)
        tf.summary.scalar('losses/high_prob_ratio', tf.reduce_mean(loss_mask)) # The ratio of unl images above the thresh
        tf.summary.scalar('losses/percent_confidence_mask', tf.reshape(curr_confidence_mask,[]))
        loss_mask = tf.stop_gradient(loss_mask)
        loss_l2u = loss_l2u * tf.expand_dims(loss_mask, axis=-1)

        # Return the average unsupervised loss.
        avg_unsup_loss = (tf.reduce_sum(loss_l2u) /
                        tf.maximum(tf.reduce_sum(loss_mask) * FLAGS.nclass, 1))
        return avg_unsup_loss 
Example #11
Source File: tf_utils.py    From transform with Apache License 2.0 5 votes vote down vote up
def _iteration_l_moments_sparse(
    current_index, l1_sum, l2_sum, l3_sum, l4_sum, count_samples,
    count_pairs, count_triplets, count_quadruplets, x_rank_2):
  """Process one column of a `SparseTensor` and updates L-moments variables."""
  current_x = tf.boolean_mask(
      x_rank_2.values,
      tf.math.equal(x_rank_2.indices[:, 1], [current_index]))
  sorted_x = tf.sort(current_x, axis=0)
  num_samples = tf.shape(current_x)[0]
  (current_samples, current_pairs, current_triplets, current_quadruplets,
   l1_factors, l2_factors, l3_factors,
   l4_factors) = _num_terms_and_factors(num_samples, x_rank_2.values.dtype)

  dim_1 = x_rank_2.dense_shape[1]
  new_l1_sum = l1_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l1_factors), axis=0)], [dim_1])
  new_l2_sum = l2_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l2_factors), axis=0)], [dim_1])
  new_l3_sum = l3_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l3_factors), axis=0)], [dim_1])
  new_l4_sum = l4_sum + tf.scatter_nd(
      [[current_index]],
      [tf.reduce_sum(tf.multiply(sorted_x, l4_factors), axis=0)], [dim_1])

  new_count_samples = count_samples + tf.scatter_nd(
      [[current_index]], [current_samples], [dim_1])
  new_count_pairs = count_pairs + tf.scatter_nd(
      [[current_index]], [current_pairs], [dim_1])
  new_count_triplets = count_triplets + tf.scatter_nd(
      [[current_index]], [current_triplets], [dim_1])
  new_count_quadruplets = count_quadruplets + tf.scatter_nd(
      [[current_index]], [current_quadruplets], [dim_1])

  return (tf.add(current_index, 1),
          new_l1_sum, new_l2_sum, new_l3_sum, new_l4_sum,
          new_count_samples, new_count_pairs, new_count_triplets,
          new_count_quadruplets, x_rank_2) 
Example #12
Source File: sample.py    From gpt-2-tensorflow2.0 with MIT License 5 votes vote down vote up
def top_p_logits(logits, p):
	"""Took from OpenAI GPT-2 Implememtation"""
	batch = tf.shape(logits)[0]
	sorted_logits = tf.sort(logits, direction='DESCENDING', axis=-1)
	cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)
	indices = tf.stack([
		tf.range(0, batch),
		tf.maximum(tf.reduce_sum(tf.cast(cumulative_probs <= p, tf.int32), axis=-1) - 1, 0),
	], axis=-1)
	min_values = tf.gather_nd(sorted_logits, indices)
	return tf.where(
		logits < min_values,
		tf.ones_like(logits) * -1e10,
		logits,
	) 
Example #13
Source File: sample.py    From galois-autocompleter with MIT License 5 votes vote down vote up
def top_p_logits(logits, p):
    with tf.variable_scope('top_p_logits'):
        logits_sort = tf.sort(logits, direction='DESCENDING')

        # logits_sort = (logits_sort[i,:] for i in range(logits_sort.get_shape().as_list()[0]))
        probs_sort = tf.nn.softmax(logits_sort)
        probs_sums = tf.cumsum(probs_sort, axis=1, exclusive=True)

        logits_masked = tf.where(probs_sums < p, logits_sort, tf.ones_like(logits_sort)*1000) # [batchsize, vocab]
        min_logits = tf.reduce_min(logits_masked, axis=1, keepdims=True) # [batchsize, 1]
        return tf.where(
            logits < min_logits,
            tf.ones_like(logits, dtype=logits.dtype) * -1e10,
            logits,
        ) 
Example #14
Source File: language_generation.py    From XLnet-gen with MIT License 4 votes vote down vote up
def sample_token(logits):
    """
    Inputs:
    logits: tf.Tensor([batch_size,len,num_tokens])
    Outpus:
    samples: tf.Tensor([batch_size,len,1])
    """
    # credits: https://github.com/nshepperd/gpt-2

    logits /= FLAGS.temperature

    batch_size = tf.shape(logits)[0]
    seq_len = tf.shape(logits)[1]
    num_toks = tf.shape(logits)[2]

    if sampling_strategy() == 'top_p':
        logits_sorted = tf.sort(logits,
                                direction="DESCENDING",
                                axis=-1)
        probs = tf.nn.softmax(logits_sorted, axis=-1)
        cum_probs = tf.math.cumsum(probs,
                                   axis=-1,
                                   exclusive=True)
        logits_masked = tf.where(cum_probs < FLAGS.top_p,
                                 logits_sorted,
                                 tf.ones_like(logits_sorted) * 100)
        min_logits = tf.reduce_min(logits_masked, axis=-1)

        logits_masked = tf.where(logits < min_logits,
                                 tf.ones_like(logits) * -1e10,
                                 logits)

    elif sampling_strategy() == "top_k":
        if FLAGS.top_k != 0:
            values, _ = tf.nn.top_k(logits, k=FLAGS.top_k)
            min_values = values[:, :, -1:]
            logits_masked = tf.where(
                logits < min_values,
                tf.ones_like(logits, dtype=logits.dtype) * -1e10,
                logits,
            )
    else:
        raise NotImplementedError("Invalid sampling strategy")

    logits_masked = tf.reshape(logits_masked, (-1, num_toks))

    samples = tf.random.categorical(logits_masked,
                                    num_samples=1,
                                    dtype=tf.int32)

    probs = tf.nn.softmax(tf.reshape(logits, (-1, num_toks)), axis=-1)
    confidences = tf.gather_nd(params=probs, batch_dims=1, indices=samples)

    return tf.reshape(samples, (batch_size, seq_len, 1)),\
        tf.reshape(confidences, (batch_size, seq_len, 1)) 
Example #15
Source File: utils_tf.py    From cleverhans with MIT License 4 votes vote down vote up
def clip_eta(eta, ord, eps):
  """
  Helper function to clip the perturbation to epsilon norm ball.
  :param eta: A tensor with the current perturbation.
  :param ord: Order of the norm (mimics Numpy).
              Possible values: np.inf, 1 or 2.
  :param eps: Epsilon, bound of the perturbation.
  """

  # Clipping perturbation eta to self.ord norm ball
  if ord not in [np.inf, 1, 2]:
    raise ValueError('ord must be np.inf, 1, or 2.')
  reduc_ind = list(xrange(1, len(eta.get_shape())))
  avoid_zero_div = 1e-12
  if ord == np.inf:
    eta = clip_by_value(eta, -eps, eps)
  elif ord == 1:
    # Implements a projection algorithm onto the l1-ball from
    # (Duchi et al. 2008) that runs in time O(d*log(d)) where d is the
    # input dimension.
    # Paper link (Duchi et al. 2008): https://dl.acm.org/citation.cfm?id=1390191

    eps = tf.cast(eps, eta.dtype)

    dim = tf.reduce_prod(tf.shape(eta)[1:])
    eta_flat = tf.reshape(eta, (-1, dim))
    abs_eta = tf.abs(eta_flat)

    if 'sort' in dir(tf):
      mu = -tf.sort(-abs_eta, axis=-1)
    else:
      # `tf.sort` is only available in TF 1.13 onwards
      mu = tf.nn.top_k(abs_eta, k=dim, sorted=True)[0]
    cumsums = tf.cumsum(mu, axis=-1)
    js = tf.cast(tf.divide(1, tf.range(1, dim + 1)), eta.dtype)
    t = tf.cast(tf.greater(mu - js * (cumsums - eps), 0), eta.dtype)

    rho = tf.argmax(t * cumsums, axis=-1)
    rho_val = tf.reduce_max(t * cumsums, axis=-1)
    theta = tf.divide(rho_val - eps, tf.cast(1 + rho, eta.dtype))

    eta_sgn = tf.sign(eta_flat)
    eta_proj = eta_sgn * tf.maximum(abs_eta - theta[:, tf.newaxis], 0)
    eta_proj = tf.reshape(eta_proj, tf.shape(eta))

    norm = tf.reduce_sum(tf.abs(eta), reduc_ind)
    eta = tf.where(tf.greater(norm, eps), eta_proj, eta)

  elif ord == 2:
    # avoid_zero_div must go inside sqrt to avoid a divide by zero
    # in the gradient through this operation
    norm = tf.sqrt(tf.maximum(avoid_zero_div,
                              reduce_sum(tf.square(eta),
                                         reduc_ind,
                                         keepdims=True)))
    # We must *clip* to within the norm ball, not *normalize* onto the
    # surface of the ball
    factor = tf.minimum(1., div(eps, norm))
    eta = eta * factor
  return eta