def sample_temperature_per_example(logits, temperature, sampling_keep_top_k=-1):
  """Either random sampling with different temperature per example.

  Args:
    logits: a Tensor.
    temperature: a float vector of same size as logits.
    sampling_keep_top_k: If not -1, only sample from the top k logits.
  Returns:
    a Tensor with one fewer dimension than logits.
  """
  logits = _select_top_k(logits, sampling_keep_top_k)
  logits /= tf.reshape(temperature, [-1] + [1] * (len(logits.shape) - 1))
  reshaped_logits = tf.reshape(logits, [-1, shape_list(logits)[-1]])
  choices = tf.multinomial(reshaped_logits, 1)
  choices = tf.reshape(choices,
                       shape_list(logits)[:logits.get_shape().ndims - 1])
  return choices
Either random sampling with different temperature per example.

  Args:
    logits: a Tensor.
    temperature: a float vector of same size as logits.
    sampling_keep_top_k: If not -1, only sample from the top k logits.
  Returns:
    a Tensor with one fewer dimension than logits.
  
tensor2tensor/layers/common_layers.py