code: def add_ema_control_dependencies(
      vector_quantizer,
      one_hot_assignments,
      codes,
      commitment_loss,
      decay,
  ):
      """Add control dependencies to the commmitment loss to update the codebook.
  
      Args:
        vector_quantizer: An instance of the VectorQuantizer class.
        one_hot_assignments: The one-hot vectors corresponding to the matched
          codebook entry for each code in the batch.
        codes: A `float`-like `Tensor` containing the latent vectors to be compared
          to the codebook.
        commitment_loss: The commitment loss from comparing the encoder outputs to
          their neighboring codebook entries.
        decay: Decay factor for exponential moving average.
  
      Returns:
        commitment_loss: Commitment loss with control dependencies.
      """
      # Use an exponential moving average to update the codebook.
      updated_ema_count = moving_averages.assign_moving_average(
          vector_quantizer.ema_count,
          tf.reduce_sum(input_tensor=one_hot_assignments, axis=[0, 1]),
          decay,
          zero_debias=False,
      )
      updated_ema_means = moving_averages.assign_moving_average(
          vector_quantizer.ema_means,
          tf.reduce_sum(
              input_tensor=tf.expand_dims(codes, 2)
              * tf.expand_dims(one_hot_assignments, 3),
              axis=[0, 1],
          ),
          decay,
          zero_debias=False,
      )
  
      # Add small value to avoid dividing by zero.
      perturbed_ema_count = updated_ema_count + 1e-5
      with tf.control_dependencies([commitment_loss]):
          update_means = tf.compat.v1.assign(
              vector_quantizer.codebook,
              updated_ema_means / perturbed_ema_count[..., tf.newaxis],
          )
          with tf.control_dependencies([update_means]):
              return tf.identity(commitment_loss)
  
imports: import functools
  import os
  import time
  from absl import flags
  from matplotlib import cm
  from matplotlib import figure
  from matplotlib.backends import backend_agg
  import numpy
  from six.moves import urllib
  import tensorflow.compat.v1
  from tensorflow_probability import distributions
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow.contrib.learn.python.learn.datasets import mnist
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 3 for calling a global function with example usage: add_ema_control_dependencies(vector_quantizer, one_hot_assignments, codes, commitment_loss, decay) and returns: tf