code: def get_vq_codebook(codebook_size, hidden_size):
    
    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=means.initialized_value(),
            trainable=False)
  
    return means, ema_means, ema_count
docs: Get lookup table for VQ bottleneck.