ยป Keras API reference / KerasNLP / Layers / PositionEmbedding layer

PositionEmbedding layer

[source]

PositionEmbedding class

keras_nlp.layers.PositionEmbedding(
    sequence_length, initializer="glorot_uniform", **kwargs
)

A layer which learns a position embedding for inputs sequences.

This class assumes that in the input tensor, the last dimension corresponds to the features, and the dimension before the last corresponds to the sequence.

This layer optionally accepts tf.RaggedTensors as inputs to process batches of sequences of different lengths. The one ragged dimension must be the dimension that corresponds to the sequence, that is, the penultimate dimension.

This layer does not supporting masking, but can be combined with a keras.layers.Embedding for padding mask support.

Arguments

  • sequence_length: The maximum length of the dynamic sequence.
  • initializer: The initializer to use for the embedding weights. Defaults to "glorot_uniform".
  • seq_axis: The axis of the input tensor where we add the embeddings.

Examples

Called directly on input.

>>> layer = keras_nlp.layers.PositionEmbedding(sequence_length=10)
>>> layer(tf.zeros((8, 10, 16))).shape
TensorShape([8, 10, 16])

Combine with a token embedding.

seq_length = 50
vocab_size = 5000
embed_dim = 128
inputs = keras.Input(shape=(seq_length,))
token_embeddings = keras.layers.Embedding(
    input_dim=vocab_size, output_dim=embed_dim
)(inputs)
position_embeddings = keras_nlp.layers.PositionEmbedding(
    sequence_length=seq_length
)(token_embeddings)
outputs = token_embeddings + position_embeddings

Reference