class_docstr: Trainable layer that maps discrete tokens/IDs to vectors.
  
  Embedding layers are commonly used to map discrete data, like words in NLP,
  into vectors. Here is a canonical example::
  
      vocab_size = 5
      word_ids = np.array([1, 2, 3, 4], dtype=np.int32)  # word_ids < vocab_size
      embedding_layer = tl.Embedding(vocab_size, 32)
      embedding_layer.init(trax.shapes.signature(word_ids))
      embedded = embedding_layer(word_ids)  # embedded.shape = (4, 32)
class_docstr_tok: ["Trainable","layer","that","maps","discrete","tokens\/IDs","to","vectors",".","Embedding","layers","are","commonly","used","to","map","discrete","data",",","like","words","in","NLP",",","into","vectors",".","Here","is","a","canonical","example",":",":","vocab_size","=","5","word_ids","=","np.array","(","[","1",",","2",",","3",",","4","]",",","dtype=np.int32",")","#","word_ids","<","vocab_size","embedding_layer","=","tl.Embedding","(","vocab_size",",","32",")","embedding_layer.init","(","trax.shapes.signature","(","word_ids",")",")","embedded","=","embedding_layer","(","word_ids",")","#","embedded.shape","=","(","4",",","32",")"]
docstr: Returns an embedding layer with given vocabulary size and vector size.
  
  The layer clips input values (token IDs) to the range `[0, vocab_size)`.
  That is, negative token IDs all clip to `0` before being mapped to a
  vector, and token IDs with value `vocab_size` or greater all clip to
  `vocab_size - 1` before being mapped to a vector.
  
  Args:
    vocab_size: Size of the input vocabulary. The layer will assign a unique
      vector to each id in `range(vocab_size)`.
    d_feature: Dimensionality/depth of the output vectors.
    use_bfloat16: If `True`, use bfloat16 weights instead of the default
      float32; this can save memory but may (rarely) lead to numerical issues.
    kernel_initializer: Function that creates (random) initial vectors for
      the embedding.
docstr_tok: ["Returns","an","embedding","layer","with","given","vocabulary","size","and","vector","size",".","The","layer","clips","input","values","(","token","IDs",")","to","the","range","`","[","0",",","vocab_size",")","`",".","That","is",",","negative","token","IDs","all","clip","to","`","0","`","before","being","mapped","to","a","vector",",","and","token","IDs","with","value","`","vocab_size","`","or","greater","all","clip","to","`","vocab_size","-","1","`","before","being","mapped","to","a","vector",".","Args",":","vocab_size",":","Size","of","the","input","vocabulary",".","The","layer","will","assign","a","unique","vector","to","each","id","in","`","range","(","vocab_size",")","`",".","d_feature",":","Dimensionality\/depth","of","the","output","vectors",".","use_bfloat16",":","If","`","True","`",",","use","bfloat16","weights","instead","of","the","default","float32",";","this","can","save","memory","but","may","(","rarely",")","lead","to","numerical","issues",".","kernel_initializer",":","Function","that","creates","(","random",")","initial","vectors","for","the","embedding","."]
code: def __init__(
      self,
      vocab_size,
      d_feature,
      use_bfloat16=False,
      kernel_initializer=init.ScaledInitializer(
          out_dim=-1,
          in_dim=-2,
          scale=1.0,
          mode="fan_out",
          distribution="uniform",
      ),
  ):
      """Returns an embedding layer with given vocabulary size and vector size.
  
      The layer clips input values (token IDs) to the range `[0, vocab_size)`.
      That is, negative token IDs all clip to `0` before being mapped to a
      vector, and token IDs with value `vocab_size` or greater all clip to
      `vocab_size - 1` before being mapped to a vector.
  
      Args:
        vocab_size: Size of the input vocabulary. The layer will assign a unique
          vector to each id in `range(vocab_size)`.
        d_feature: Dimensionality/depth of the output vectors.
        use_bfloat16: If `True`, use bfloat16 weights instead of the default
          float32; this can save memory but may (rarely) lead to numerical issues.
        kernel_initializer: Function that creates (random) initial vectors for
          the embedding.
      """
      # TODO(jonni): is the clipping behavior what we want going forward?
      super().__init__(name=f"Embedding_{vocab_size}_{d_feature}")
      self._d_feature = d_feature  # feature dimensionality
      self._vocab_size = vocab_size
      self._use_bfloat16 = use_bfloat16
      self._kernel_initializer = kernel_initializer
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 1 to create a new Embedding object from inherited base classes: base with example: obj = Embedding(vocab_size, d_feature, use_bfloat16, kernel_initializer)