code: def multinomial_sample(x, vocab_size=None, sampling_method="random",
                         temperature=1.0):
    
    vocab_size = vocab_size or common_layers.shape_list(x)[-1]
    if sampling_method == "random" and temperature > 0.0:
      samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
    else:
      samples = tf.argmax(x, axis=-1)
    reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
    return reshaped_samples
docs: Multinomial sampling from a n-dimensional tensor.
  
    Args:
      x: Tensor of shape [..., vocab_size]. Parameterizes logits of multinomial.
      vocab_size: Number of classes in multinomial distribution.
      sampling_method: String, "random" or otherwise deterministic.
      temperature: Positive float.
  
    Returns:
      Tensor of shape [...].