code: def gumbel_softmax(
      x,
      z_size,
      mode,
      softmax_k=0,
      temperature_warmup_steps=150000,
      summary=True,
      name=None,
  ):
      """Gumbel softmax discretization bottleneck.
  
      Args:
        x: Input to the discretization bottleneck.
        z_size: Number of bits, where discrete codes range from 1 to 2**z_size.
        mode: tf.estimator.ModeKeys.
        softmax_k: If > 0 then do top-k softmax.
        temperature_warmup_steps: Number of steps it takes to decay temperature to
          0.
        summary: Whether to write summaries.
        name: Name for the bottleneck scope.
  
      Returns:
        Embedding function, discrete code, and loss.
      """
      with tf.variable_scope(name, default_name="gumbel_softmax"):
          m = tf.layers.dense(x, 2**z_size, name="mask")
          if softmax_k > 0:
              m, kl = top_k_softmax(m, softmax_k)
              return m, m, 1.0 - tf.reduce_mean(kl)
          logsm = tf.nn.log_softmax(m)
  
          # Gumbel-softmax sample.
          gumbel_samples = gumbel_sample(common_layers.shape_list(m))
          steps = temperature_warmup_steps
          gumbel_samples *= (
              common_layers.inverse_exp_decay(steps // 5) * 0.5
          )
          temperature = 1.2 - common_layers.inverse_lin_decay(steps)
  
          # 10% of the time keep reasonably high temperature to keep learning.
          temperature = tf.cond(
              tf.less(tf.random_uniform([]), 0.9),
              lambda: temperature,
              lambda: tf.random_uniform([], minval=0.5, maxval=1.0),
          )
          s = tf.nn.softmax((logsm + gumbel_samples) / temperature)
          m = tf.nn.softmax(m)
          kl = -tf.reduce_max(logsm, axis=-1)
  
          if summary:
              tf.summary.histogram("max-log", tf.reshape(kl, [-1]))
  
          # Calculate the argmax and construct hot vectors.
          maxvec = tf.reshape(tf.argmax(m, axis=-1), [-1])
          maxvhot = tf.stop_gradient(tf.one_hot(maxvec, 2**z_size))
  
          # Add losses that prevent too few being used.
          distrib = tf.reshape(logsm, [-1, 2**z_size]) * maxvhot
          d_mean = tf.reduce_mean(distrib, axis=[0], keep_dims=True)
          d_variance = tf.reduce_mean(
              tf.squared_difference(distrib, d_mean), axis=[0]
          )
          d_dev = -tf.reduce_mean(d_variance)
          ret = s
  
          if mode != tf_estimator.ModeKeys.TRAIN:
              ret = tf.reshape(
                  maxvhot, common_layers.shape_list(s)
              )  # Just hot @eval.
          return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 12 for calling a global function with example usage: gumbel_softmax(x, z_size, mode, softmax_k, temperature_warmup_steps, summary, name) and returns: m,  ret,  unknown, m,  m,  unknown