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 embedding vectors corresponding to input token IDs.
  
  Args:
    x: Tensor of token IDs.
  
  Returns:
    Tensor of embedding vectors.
docstr_tok: ["Returns","embedding","vectors","corresponding","to","input","token","IDs",".","Args",":","x",":","Tensor","of","token","IDs",".","Returns",":","Tensor","of","embedding","vectors","."]
code: def forward(self, x):
      """Returns embedding vectors corresponding to input token IDs.
  
      Args:
        x: Tensor of token IDs.
  
      Returns:
        Tensor of embedding vectors.
      """
      embedded = jnp.take(self.weights, x, axis=0, mode="clip")
      if (
          self._use_bfloat16
      ):  # Return float32 activations w/ bfloat16 weights.
          embedded = embedded.astype(jnp.float32)
      return embedded
  
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 2 for calling the Embedding obj's underlying member method code with example usage: obj.forward(x) and returns: embedded