diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py index df46b3622b..d002d8745f 100644 --- a/keras_nlp/layers/token_and_position_embedding.py +++ b/keras_nlp/layers/token_and_position_embedding.py @@ -23,9 +23,12 @@ @keras.utils.register_keras_serializable(package="keras_nlp") class TokenAndPositionEmbedding(keras.layers.Layer): """A layer which sums a token and position embedding. - - This layer assumes that the last dimension in the input corresponds - to the sequence dimension. + + Token and position embeddings are ways of representing words and their order + in a sentence. This layer creates a `keras.layers.Embedding` token embedding + and a `keras_nlp.layers.PositionEmbedding` position embedding and sums their + output when called. This layer assumes that the last dimension in the input + corresponds to the sequence dimension. Args: vocabulary_size: The size of the vocabulary.