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
Returns embedding vectors corresponding to input token IDs.

    Args:
      x: Tensor of token IDs.

    Returns:
      Tensor of embedding vectors.
    
trax/layers/core.py