def _select_top_k(logits, top_k):
  """Replaces logits, expect the top k highest values, with small number (-1e6).

  If k is -1 don't replace anything.

  Args:
    logits: A `Tensor` of shape [batch_size, ..., vocab_size]
    top_k: vector of batch size.

  Returns:
    A `Tensor` with same shape  as logits.
  """
  vocab_size = logits.shape[-1]

  top_k = tf.where(
      tf.not_equal(top_k, -1), top_k,
      tf.ones_like(top_k) * vocab_size)

  return tf.where(
      tf.argsort(logits) < tf.reshape(top_k, [-1] + [1] *
                                      (len(logits.shape) - 1)), logits,
      tf.ones_like(logits) * -1e6)
Replaces logits, expect the top k highest values, with small number (-1e6).

  If k is -1 don't replace anything.

  Args:
    logits: A `Tensor` of shape [batch_size, ..., vocab_size]
    top_k: vector of batch size.

  Returns:
    A `Tensor` with same shape  as logits.
  
tensor2tensor/layers/common_layers.py