code: def vq_discrete_bottleneck(x, hparams):
      """Simple vector quantized discrete bottleneck."""
      tf.logging.info("Using EMA with beta = {}".format(hparams.beta))
      bottleneck_size = 2**hparams.bottleneck_bits
      x_shape = common_layers.shape_list(x)
      x = tf.reshape(x, [-1, hparams.hidden_size])
      x_means_hot, e_loss = vq_nearest_neighbor(x, hparams)
      means, ema_means, ema_count = (
          hparams.means,
          hparams.ema_means,
          hparams.ema_count,
      )
  
      # Update the ema variables
      updated_ema_count = moving_averages.assign_moving_average(
          ema_count,
          tf.reduce_sum(x_means_hot, axis=0),
          hparams.decay,
          zero_debias=False,
      )
  
      dw = tf.matmul(x_means_hot, x, transpose_a=True)
      updated_ema_means = moving_averages.assign_moving_average(
          ema_means, dw, hparams.decay, zero_debias=False
      )
      n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
      updated_ema_count = (
          (updated_ema_count + hparams.epsilon)
          / (n + bottleneck_size * hparams.epsilon)
          * n
      )
      updated_ema_means = updated_ema_means / tf.expand_dims(
          updated_ema_count, axis=-1
      )
      with tf.control_dependencies([e_loss]):
          update_means = tf.assign(means, updated_ema_means)
          with tf.control_dependencies([update_means]):
              loss = hparams.beta * e_loss
  
      discrete = tf.reshape(
          x_means_hot, x_shape[:-1] + [bottleneck_size]
      )
      return discrete, loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from six.moves import range
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import beam_search
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 3 for calling a global function with example usage: vq_discrete_bottleneck(x, hparams) and returns: discrete,  loss