class_docstr: Tests for discretization layers.
code: def testGetVQBottleneck(self):
      bottleneck_bits = 2
      bottleneck_size = 2**bottleneck_bits
      hidden_size = 3
      means, _, ema_count = discretization.get_vq_codebook(
          bottleneck_size, hidden_size
      )
      assign_op = means.assign(
          tf.zeros(shape=[bottleneck_size, hidden_size])
      )
      means_new, _, _ = discretization.get_vq_codebook(
          bottleneck_size, hidden_size
      )
      with self.test_session() as sess:
          tf.global_variables_initializer().run()
          sess.run(assign_op)
          self.assertTrue(np.all(sess.run(means_new) == 0))
          self.assertTrue(np.all(sess.run(ema_count) == 0))
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.layers import discretization
  from tensor2tensor.utils import test_utils
  import tensorflow.compat.v1
label_desc: Use image node_id 10 for calling the DiscretizationTest obj's underlying member method code with example usage: obj.testGetVQBottleneck() without return types