code: def multinomial_sample(
      x, vocab_size=None, sampling_method="random", temperature=1.0
  ):
      """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 [...].
      """
      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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import transformer_layers
  from tensor2tensor.utils import beam_search
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 3 for calling a global function with example usage: multinomial_sample(x, vocab_size, sampling_method, temperature) and returns: reshaped_samples