Easy Primitives

Implement Embedding Lookup

An embedding table stores one vector per item, such as a token in a vocabulary. An integer ID selects a row of that table. A lookup retrieves the stored numbers; it does not compute a new vector from the ID itself.

For the table [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]] and IDs [2, 0, 2], return [[0.5, 0.6], [0.1, 0.2], [0.5, 0.6]]. Preserve the requested order and repeat a row whenever its ID appears again.

Your task

Implement embedding_lookup(weight, indices):

weight is a floating-point table of shape (V, D), with positive vocabulary size V and embedding dimension D.
indices is a 1D integer tensor of shape (S,). Every ID is in [0, V - 1]; no handling for negative or out-of-range IDs is required.
Return a tensor of shape (S, D), with output row i equal to weight[indices[i]].

Use tensor indexing rather than a built-in embedding layer. Do not sort, deduplicate, sum, or average the selected vectors. Keep both output axes, even when S=1 or D=1. An empty index vector returns shape (0, D). The supplied table is already populated; you do not initialize or train it.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Look up a learned vector for each token

An embedding table E has shape (vocabulary_size, embedding_dim). Integer token ID i selects row E_i. IDs are labels for rows, not measurements: token 100 does not have twice the meaning of token 50.

A lookup is equivalent to multiplying a one-hot token vector by E, but avoids constructing the mostly zero vector. Looking up IDs [2, 2, 5] returns rows [E_2, E_2, E_5]. Both uses of row 2 contribute gradients to that row during training, subject to any scaling or special padding behavior in the task.

Averaging token embeddings gives one fixed-width vector for a sequence. It loses order information: the average for tokens [2, 5] is the same as for [5, 2].

Reuse the embedding weights to predict tokens

Weight tying uses the same table for input embeddings and output scores. For a hidden vector h with width embedding_dim, the product hE^T gives one score, or logit, for every vocabulary entry.

An input lookup selects one row. The output projection compares the hidden vector with every row. Sharing the table saves a separate output weight matrix, but means both uses must contribute gradients to the same parameter. Two copies that are updated independently are no longer tied.

Check that token IDs are in range and have the required integer dtype. When pooling padded sequences, exclude padding from both the sum and the count. Even a zero padding vector lowers an ordinary average if it is included in the denominator. Follow the specified behavior for a sequence with no valid tokens.

Detecting runtime…

Implement Embedding Lookup

Your task

Look up a learned vector for each token

Reuse the embedding weights to predict tokens

Your code, saved here

Stuck?

PyTorch reference solution