Medium End-to-End

Word Embedding Model

An embedding table stores one vector per vocabulary item. To represent a sequence with a single vector, look up its embeddings and average them, coordinate by coordinate. This is called mean pooling.

Your task

Implement embed_and_average(embedding_table, indices).

embedding_table is a floating-point tensor of shape (vocab_size, embed_dim).
indices has shape (seq_len,), with seq_len >= 1. Each value is a valid vocabulary ID in 0 <= id < vocab_size. The runtime supplies integer-valued floats, so convert them to an integer dtype before indexing.
Return a floating-point vector of shape (embed_dim,), preserving the table's dtype. Leave both inputs unchanged.

output[d] = sum(embedding_table[indices[t], d] for t in the sequence) / seq_len

Average over sequence positions, not embedding coordinates or all vocabulary rows. Repeated IDs contribute once per occurrence. There is no padding mask or special ID to skip: every supplied position counts in the denominator.

For example, with table [[2, -1], [8, 5]] and IDs [0, 0, 1], average [2, -1], [2, -1] and [8, 5] to get [4, 1]. For one ID, return that row as a vector, without a leading size-1 dimension.

The table is supplied, so this exercise does not train embeddings, tokenize text or apply a classifier. Mean pooling also discards word order: rearranging the same IDs leaves the mathematical average unchanged. Sequences with different word orders can therefore receive the same representation.

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…

Word Embedding Model

Your task

Look up a learned vector for each token

Reuse the embedding weights to predict tokens

Your code, saved here

Stuck?

JAX reference solution