Token embedding lookup
The problem, from first principles
[01-tokenization/05-encode-decode-roundtrip]'s encode produces a sequence of integer token ids, but an integer id is a completely ARBITRARY label: id 47 and id 48 might refer to two totally unrelated words, and there's no reason the network should treat them as "close" just because their id numbers happen to be adjacent. What a network actually needs is a DENSE VECTOR representation of each token, one where semantically similar words end up with similar vectors (learned automatically during training, not hand-designed), so [03-dl-training/02-layers/01-linear-forward]'s matrix multiplications and everything built on top of them have something meaningful to compute with.
The embedding table is exactly this: one row per vocabulary entry, each row a learnable vector of embed_dim numbers. "Looking up" a token's embedding is conceptually just a lookup table access (in exactly the sense the "lookup table" comparison in [03-dl-training/05-why-deep-networks-work/02-representation-learning] described, except here the vocabulary size is fixed and small enough that a genuine per-token row DOES make sense, unlike that question's exponentially-many feature COMBINATIONS), but framed as a matrix operation so it fits naturally into a network trained end-to-end with [02-deep-learning-core/04-autograd]'s backpropagation.
From theory to code
Implement embedding_forward(token_ids, embedding_table). embedding_table has shape (vocab_size, embed_dim), one row per vocabulary entry. token_ids can be any shape, a single sequence of ids, or a whole batch of sequences. The output replaces every id in token_ids with that id's corresponding row from embedding_table, so the output shape is token_ids's shape with one extra trailing embed_dim dimension.
Constraints
- For a single id,
embedding_forwardreturns exactlyembedding_table[id], that row, unchanged. - For
token_idsof shape(batch_size, seq_len), the output must have shape(batch_size, seq_len, embed_dim). - The SAME id appearing multiple times in
token_idsmust return the SAME row every time (the embedding table doesn't change during a single forward pass). - Must not modify
embedding_tableitself.
Hints
Hint 1
NumPy's fancy indexing already does exactly this in one line: embedding_table[token_ids]. If token_ids has shape (batch_size, seq_len) and embedding_table has shape (vocab_size, embed_dim), the result automatically comes out as (batch_size, seq_len, embed_dim).
Hint 2
Fancy indexing with an array of ids (rather than a single integer or a slice) always returns a NEW array (a copy of the selected rows), so embedding_table itself is never modified by this operation, exactly matching the constraint above.
Click "Run Tests" to test your implementation