def __init__(self,
               vocab_size,
               d_feature,
               use_bfloat16=False,
               kernel_initializer=init.ScaledInitializer(
                   out_dim=-1, in_dim=-2, scale=1., 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
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.
    
trax/layers/core.py