Skip to content

← Tokenization & Embeddings step 5 of 15

Easy Primitives

Learned Absolute Position Embedding

A learned position embedding table stores one vector for each absolute position: row 0 represents position 0, row 1 represents position 1, and so on. A transformer can add these vectors to token embeddings to supply position information. During training, the table values can be updated by backpropagation.

This exercise implements only the lookup. The table is supplied; do not initialize or train it, or add the result to token embeddings.

Your task

Implement learned_position_embedding(positions, embedding_table).

  • positions has shape (N, T): one position ID for each token in each batch row. The runtime delivers integer-valued floats, so cast them to an integer dtype before indexing. Every ID is in the range 0 <= position < max_len.
  • embedding_table has shape (max_len, d_model) and is shared by the whole batch.
  • Return shape (N, T, d_model), preserving the table's floating-point dtype.
output[n, t, :] = embedding_table[positions[n, t], :]

Use the supplied IDs, even when they repeat, skip positions or appear in reverse order. Do not replace them with 0, 1, ..., T-1. Repeated IDs select the same vector, including when they occur in different batch rows.

For example, with positions=[[2, 0, 2]] and the table [[1, 2], [3, 4], [5, 6]], return [[[5, 6], [1, 2], [5, 6]]].

What makes the positions absolute?

Each ID names a position directly, rather than a distance between two tokens. The supplied table only has rows for positions 0 through max_len-1. Looking up a longer sequence requires valid IDs within that range; this function does not extend the table.

The illustration below starts on the Learned tab. Its values are synthetic, not trained weights or the table used by the tests.

Loading visualization…