code: def sample_with_temperature(logits, temperature):
      """Either argmax after softmax or random sample along the pitch axis.
  
      Args:
        logits: a Tensor of shape (batch, time, pitch, instrument).
        temperature: a float  0.0=argmax 1.0=random
  
      Returns:
        a Tensor of the same shape, with one_hots on the pitch dimension.
      """
      logits = tf.transpose(logits, [0, 1, 3, 2])
      pitch_range = tf.shape(logits)[-1]
  
      def sample_from_logits(logits):
          with tf.control_dependencies(
              [tf.assert_greater(temperature, 0.0)]
          ):
              logits = tf.identity(logits)
          reshaped_logits = (
              tf.reshape(logits, [-1, tf.shape(logits)[-1]])
              / temperature
          )
          choices = tf.multinomial(reshaped_logits, 1)
          choices = tf.reshape(
              choices, tf.shape(logits)[: logits.get_shape().ndims - 1]
          )
          return choices
  
      choices = tf.cond(
          tf.equal(temperature, 0.0),
          lambda: tf.argmax(tf.nn.softmax(logits), -1),
          lambda: sample_from_logits(logits),
      )
      samples_onehot = tf.one_hot(choices, pitch_range)
      return tf.transpose(samples_onehot, [0, 1, 3, 2])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import time
  from magenta.models.coconet import lib_graph
  from magenta.models.coconet import lib_hparams
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling a global function with example usage: sample_with_temperature(logits, temperature) and returns: tf, choices