code: def pixels_from_softmax(
      frame_logits,
      pure_sampling=False,
      temperature=1.0,
      gumbel_noise_factor=0.2,
  ):
      """Given frame_logits from a per-pixel softmax, generate colors."""
      # If we're purely sampling, just sample each pixel.
      if pure_sampling or temperature == 0.0:
          return common_layers.sample_with_temperature(
              frame_logits, temperature
          )
  
      # Gumbel-sample from the pixel sofmax and average by pixel values.
      pixel_range = tf.to_float(tf.range(256))
      for _ in range(len(frame_logits.get_shape().as_list()) - 1):
          pixel_range = tf.expand_dims(pixel_range, axis=0)
  
      frame_logits = tf.nn.log_softmax(frame_logits)
      gumbel_samples = (
          discretization.gumbel_sample(
              common_layers.shape_list(frame_logits)
          )
          * gumbel_noise_factor
      )
  
      frame = tf.nn.softmax(
          (frame_logits + gumbel_samples) / temperature, axis=-1
      )
      result = tf.reduce_sum(frame * pixel_range, axis=-1)
      # Round on the forward pass, not on the backward one.
      return result + tf.stop_gradient(tf.round(result) - result)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import six
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.layers import discretization
  from tensor2tensor.layers import modalities
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: pixels_from_softmax(frame_logits, pure_sampling, temperature, gumbel_noise_factor) and returns: unknown, common_layers