code: def get_vq_codebook(codebook_size, hidden_size):
      """Get lookup table for VQ bottleneck."""
      with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
          means = tf.get_variable(
              name="means",
              shape=[codebook_size, hidden_size],
              initializer=tf.uniform_unit_scaling_initializer(),
          )
  
          ema_count = tf.get_variable(
              name="ema_count",
              shape=[codebook_size],
              initializer=tf.constant_initializer(0),
              trainable=False,
          )
  
          with tf.colocate_with(means):
              ema_means = tf.get_variable(
                  name="ema_means",
                  initializer=tf.cond(
                      tf.is_variable_initialized(means),
                      means.read_value,
                      lambda: means.initial_value,
                  ),
                  trainable=False,
              )
  
      return means, ema_means, ema_count
  
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 15 for calling a global function with example usage: get_vq_codebook(codebook_size, hidden_size) and returns: means,  ema_means,  ema_count