docs: def call(self, x):
      """Get token embeddings of x.
  
      Args:
        x: An int64 tensor with shape [batch_size, length]
      Returns:
        embeddings: float32 tensor with shape [batch_size, length, embedding_size]
        padding: float32 tensor with shape [batch_size, length] indicating the
          locations of the padding tokens in x.
      """
      with tf.name_scope("embedding"):
        embeddings = tf.gather(self.shared_weights, x)
  
        # Scale embedding by the sqrt of the hidden size
        embeddings *= self.hidden_size ** 0.5
  
        # Create binary array of size [batch_size, length]
        # where 1 = padding, 0 = not padding
        padding = model_utils.get_padding(x)
  
        # Set all padding embedding values to 0
        embeddings *= tf.expand_dims(1 - padding, -1)
        return embeddings