class_docstr: Performs an embedding lookup suitable for accelerator devices.
  
  This layer uses either tf.gather or tf.one_hot to translate integer indices to
  float embeddings.
  
  Attributes:
    vocab_size: Number of elements in the vocabulary.
    embedding_width: Output size of the embedding layer.
    initializer: The initializer to use for the embedding weights. Defaults to
      "glorot_uniform".
    use_one_hot: Whether to use tf.one_hot over tf.gather for the embedding
      lookup. Defaults to False (that is, using tf.gather). Setting this option
      to True may improve performance, especially on small vocabulary sizes,
      but will generally require more memory.
code: def get_config(self):
      config = {
          "vocab_size": self._vocab_size,
          "embedding_width": self._embedding_width,
          "initializer": self._initializer,
          "use_one_hot": self._use_one_hot,
      }
      base_config = super(OnDeviceEmbedding, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  
file_path: autodist/examples/benchmark/utils/modeling/layers/on_device_embedding.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from utils import tf_utils
label_desc: Use image node_id 2 for calling the OnDeviceEmbedding obj's underlying member method code with example usage: obj.get_config() and returns: dict