(logits_BxN, top_k=0, top_p=0.0, temperature=0.0)
return logits_BxN
Process logits using gumbel noise and mask top_k or top_p.

  The downstream task can perform probability sampling using gumbel-max trick
  (taking the argmax of processed logits) (Statistical theory of extreme values
  and some practical applications: a series of lectures. 1954).
  Use cases:
    greedy: top_k=0, top_p=0.0, temperature=0.0
    random sampling: top_k=0, top_p=0.0, temperature=1.0
    topk sampling: top_k=k, top_p=0.0, temperature=1.0
    nucleus sampling: top_k=0, top_p=p, temperature=1.0
    random sampling biased toward greedy: top_k=0, top_p=0.0, temperature=0.5
  Notations:
    B: batch_size, N: number of logits, K: topk value.
  Args:
    logits_BxN: tensor of [batch_size vocab_size]
    top_k: k in top_k sampling.
    top_p: probability in necleus sampling.
    temperature: gumbel noise sampling temperature.

  Returns:
    logits: processed logits which is original logits add gumbel noise and
    values outside top_k and top_p set to -inf.
Process logits using gumbel noise and mask top_k or top_p.
Process logits using gumbel noise and mask top_k or top_p .
def process_logits(logits_BxN, top_k=0, top_p=0.0, temperature=0.0):
  """Process logits using gumbel noise and mask top_k or top_p.

  The downstream task can perform probability sampling using gumbel-max trick
  (taking the argmax of processed logits) (Statistical theory of extreme values
  and some practical applications: a series of lectures. 1954).
  Use cases:
    greedy: top_k=0, top_p=0.0, temperature=0.0
    random sampling: top_k=0, top_p=0.0, temperature=1.0
    topk sampling: top_k=k, top_p=0.0, temperature=1.0
    nucleus sampling: top_k=0, top_p=p, temperature=1.0
    random sampling biased toward greedy: top_k=0, top_p=0.0, temperature=0.5
  Notations:
    B: batch_size, N: number of logits, K: topk value.
  Args:
    logits_BxN: tensor of [batch_size vocab_size]
    top_k: k in top_k sampling.
    top_p: probability in necleus sampling.
    temperature: gumbel noise sampling temperature.

  Returns:
    logits: processed logits which is original logits add gumbel noise and
    values outside top_k and top_p set to -inf.
  """
  if top_k > 0 and top_p > 0:
    raise ValueError(
        "Only one of the top_k and nucleus sampling should be specified.")

  if top_k > 0:
    top_values_BxK, _ = tf.math.top_k(logits_BxN, k=top_k, sorted=False)
    min_value_Bx1 = tf.reduce_min(top_values_BxK, axis=-1, keepdims=True)
    mask_BxN = tf.cast(tf.less(logits_BxN, min_value_Bx1), logits_BxN.dtype)
    logits_BxN -= mask_BxN * logits_BxN.dtype.max

  if top_p > 0:
    sort_indices_BxN = tf.argsort(logits_BxN, axis=-1, direction="DESCENDING")
    probs_BxN = tf.gather(
        tf.nn.softmax(logits_BxN), sort_indices_BxN, batch_dims=1)
    cumprobs_BxN = tf.cumsum(probs_BxN, axis=-1, exclusive=True)
    # The top 1 candidate always will not be masked.
    # This way ensures at least 1 indices will be selected.
    sort_mask_BxN = tf.cast(tf.greater(cumprobs_BxN, top_p), logits_BxN.dtype)
    batch_indices_BxN = tf.tile(
        tf.expand_dims(tf.range(logits_BxN.shape[0]), axis=-1),
        [1, logits_BxN.shape[1]])
    top_p_mask_BxN = tf.scatter_nd(
        tf.stack([batch_indices_BxN, sort_indices_BxN], axis=-1), sort_mask_BxN,
        logits_BxN.shape)
    logits_BxN -= top_p_mask_BxN * logits_BxN.dtype.max

  if temperature > 0:
    logits_shape = tf.shape(logits_BxN)
    uniform_noise_BxN = tf.random_uniform(logits_shape)
    logits_BxN += -tf.log(-tf.log(uniform_noise_BxN)) * temperature
  return logits_BxN
def process_logits ( logits_BxN , top_k = 0 , top_p = 0.0 , temperature = 0.0 ) : if top_k > 0 and top_p > 0 : raise ValueError ( "Only one of the top_k and nucleus sampling should be specified." ) if top_k > 0 : top_values_BxK , _ = tf . math . top_k ( logits_BxN , k = top_k , sorted = False ) min_value_Bx1 = tf . reduce_min ( top_values_BxK , axis = - 1 , keepdims = True ) mask_BxN = tf . cast ( tf . less ( logits_BxN , min_value_Bx1 ) , logits_BxN . dtype ) logits_BxN -= mask_BxN * logits_BxN . dtype . max if top_p > 0 : sort_indices_BxN = tf . argsort ( logits_BxN , axis = - 1 , direction = "DESCENDING" ) probs_BxN = tf . gather ( tf . nn . softmax ( logits_BxN ) , sort_indices_BxN , batch_dims = 1 ) cumprobs_BxN = tf . cumsum ( probs_BxN , axis = - 1 , exclusive = True ) # The top 1 candidate always will not be masked. # This way ensures at least 1 indices will be selected. sort_mask_BxN = tf . cast ( tf . greater ( cumprobs_BxN , top_p ) , logits_BxN . dtype ) batch_indices_BxN = tf . tile ( tf . expand_dims ( tf . range ( logits_BxN . shape [ 0 ] ) , axis = - 1 ) , [ 1 , logits_BxN . shape [ 1 ] ] ) top_p_mask_BxN = tf . scatter_nd ( tf . stack ( [ batch_indices_BxN , sort_indices_BxN ] , axis = - 1 ) , sort_mask_BxN , logits_BxN . shape ) logits_BxN -= top_p_mask_BxN * logits_BxN . dtype . max if temperature > 0 : logits_shape = tf . shape ( logits_BxN ) uniform_noise_BxN = tf . random_uniform ( logits_shape ) logits_BxN += - tf . log ( - tf . log ( uniform_noise_BxN ) ) * temperature return logits_BxN