code: def vq_body(
      x,
      codebook_size,
      beta=0.25,
      decay=0.999,
      epsilon=1e-5,
      soft_em=False,
      num_samples=10,
      temperature=None,
      do_update=True,
  ):
      """Discretize each x into one of codebook_size codes."""
      x_shape = common_layers.shape_list(x)
      hidden_size = x_shape[-1]
      means, ema_means, ema_count = get_vq_codebook(
          codebook_size, hidden_size
      )
      x = tf.reshape(x, [-1, hidden_size])
      x_means_hot, e_loss, distances = vq_nearest_neighbor(
          x,
          means,
          soft_em=soft_em,
          num_samples=num_samples,
          temperature=temperature,
      )
  
      def loss_with_update():
          """Update the ema variables and return loss triggering the update."""
          updated_ema_count = moving_averages.assign_moving_average(
              ema_count,
              tf.reduce_sum(
                  tf.reshape(x_means_hot, shape=[-1, codebook_size]),
                  axis=0,
              ),
              decay,
              zero_debias=False,
          )
  
          dw = tf.matmul(x_means_hot, x, transpose_a=True)
          updated_ema_means = tf.identity(
              moving_averages.assign_moving_average(
                  ema_means, dw, decay, zero_debias=False
              )
          )
          n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
          updated_ema_count = (
              (updated_ema_count + epsilon)
              / (n + codebook_size * epsilon)
              * n
          )
          updated_ema_means /= tf.expand_dims(
              updated_ema_count, axis=-1
          )
          with tf.control_dependencies([e_loss]):
              update_means = means.assign(updated_ema_means)
              with tf.control_dependencies([update_means]):
                  return beta * e_loss
  
      # Loss, also do update if requested.
      if do_update:
          loss = loss_with_update()
      else:
          loss = tf.cond(
              do_update, loss_with_update, lambda: beta * e_loss
          )
  
      d = tf.reshape(x_means_hot, x_shape[:-1] + [codebook_size])
      return d, loss, distances
  
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 18 for calling a global function with example usage: vq_body(x, codebook_size, beta, decay, epsilon, soft_em, num_samples, temperature, do_update) and returns: d,  loss,  distances, unknown