class_docstr: Creates a vector-quantizer.
  
  It quantizes a continuous vector under a codebook. The codebook is also known
  as "embeddings" or "memory", and it is learned using an exponential moving
  average.
class_docstr_tok: ["Creates","a","vector-quantizer",".","It","quantizes","a","continuous","vector","under","a","codebook",".","The","codebook","is","also","known","as","``","embeddings","''","or","``","memory","''",",","and","it","is","learned","using","an","exponential","moving","average","."]
docstr: Uses codebook to find nearest neighbor for each code.
  
  Args:
    codes: A `float`-like `Tensor` containing the latent
      vectors to be compared to the codebook. These are rank-3 with shape
      `[batch_size, latent_size, code_size]`.
  
  Returns:
    nearest_codebook_entries: The 1-nearest neighbor in Euclidean distance for
      each code in the batch.
    one_hot_assignments: The one-hot vectors corresponding to the matched
      codebook entry for each code in the batch.
docstr_tok: ["Uses","codebook","to","find","nearest","neighbor","for","each","code",".","Args",":","codes",":","A","`","float","`","-like","`","Tensor","`","containing","the","latent","vectors","to","be","compared","to","the","codebook",".","These","are","rank-3","with","shape","`","[","batch_size",",","latent_size",",","code_size","]","`",".","Returns",":","nearest_codebook_entries",":","The","1-nearest","neighbor","in","Euclidean","distance","for","each","code","in","the","batch",".","one_hot_assignments",":","The","one-hot","vectors","corresponding","to","the","matched","codebook","entry","for","each","code","in","the","batch","."]
code: def __call__(self, codes):
      """Uses codebook to find nearest neighbor for each code.
  
      Args:
        codes: A `float`-like `Tensor` containing the latent
          vectors to be compared to the codebook. These are rank-3 with shape
          `[batch_size, latent_size, code_size]`.
  
      Returns:
        nearest_codebook_entries: The 1-nearest neighbor in Euclidean distance for
          each code in the batch.
        one_hot_assignments: The one-hot vectors corresponding to the matched
          codebook entry for each code in the batch.
      """
      distances = tf.norm(
          tensor=tf.expand_dims(codes, 2)
          - tf.reshape(
              self.codebook, [1, 1, self.num_codes, self.code_size]
          ),
          axis=3,
      )
      assignments = tf.argmin(input=distances, axis=2)
      one_hot_assignments = tf.one_hot(
          assignments, depth=self.num_codes
      )
      nearest_codebook_entries = tf.reduce_sum(
          input_tensor=tf.expand_dims(one_hot_assignments, -1)
          * tf.reshape(
              self.codebook, [1, 1, self.num_codes, self.code_size]
          ),
          axis=2,
      )
      return nearest_codebook_entries, one_hot_assignments
  
imports: import functools
  import os
  import time
  from absl import flags
  from matplotlib import cm
  from matplotlib import figure
  from matplotlib.backends import backend_agg
  import numpy
  from six.moves import urllib
  import tensorflow.compat.v1
  from tensorflow_probability import distributions
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow.contrib.learn.python.learn.datasets import mnist
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 2 for calling the VectorQuantizer obj's underlying member method code with example usage: obj.__call__(codes) and returns: nearest_codebook_entries,  one_hot_assignments