code: def pixels_from_softmax(frame_logits, pure_sampling=False,
                          temperature=1.0, gumbel_noise_factor=0.2):
    
    # 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)
docs: Given frame_logits from a per-pixel softmax, generate colors.