code: def init_vq_bottleneck(bottleneck_size, hidden_size):
      """Get lookup table for VQ bottleneck."""
      means = tf.get_variable(
          name="means",
          shape=[bottleneck_size, hidden_size],
          initializer=tf.uniform_unit_scaling_initializer(),
      )
      ema_count = tf.get_variable(
          name="ema_count",
          shape=[bottleneck_size],
          initializer=tf.constant_initializer(0),
          trainable=False,
      )
      with tf.colocate_with(means):
          ema_means = tf.get_variable(
              name="ema_means",
              initializer=means.initialized_value(),
              trainable=False,
          )
  
      return means, ema_means, ema_count
  
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 1 for calling a global function with example usage: init_vq_bottleneck(bottleneck_size, hidden_size) and returns: means,  ema_means,  ema_count