Embeddings · L3 Code

Implementing an embedding lookup and a mini word2vec-style training loop, then computing nearest neighbors.

python

import numpy as np

vocab_size, embed_dim = 1000, 16
rng = np.random.default_rng(0)
embedding_table = rng.normal(0, 0.1, size=(vocab_size, embed_dim))  # learned parameters

def lookup(token_ids, table):
    return table[token_ids]     # fancy indexing: (n_tokens,) -> (n_tokens, embed_dim)

token_ids = np.array([42, 7, 999])
vectors = lookup(token_ids, embedding_table)
print(vectors.shape)   # (3, 16)

An embedding lookup table implemented as plain array indexing.

This is literally what nn.Embedding does under the hood — table[token_ids] — with the added machinery of tracking gradients so the table's rows update during training.