Token + position embeddings·8 min

Token embeddings

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage token-embeddings --watch

Start with just the token half: map each ID to its row in wte. No math, just a gather — you copy rows out of a table, nothing more. For id 0, that first row starts [-0.1101, -0.0393, 0.0331, …] — the exact floats you read straight off disk in Chapter 1.

Gather: one row of wte per id

inputIds[T] ──gather rows of wte[50257, 768]──▶ [T, 768]

What to build

Your task

Implement runTokenEmbed

Implement runTokenEmbed(inputIds): for each id, return row wte[id], giving [T, 768]. Load wte with your stage-1 loader (THEINFERENCE_WEIGHTS / ./weights/model.safetensors).

tests/adapters.ts

export function runTokenEmbed(inputIds: number[]): number[][] {
  // -> [T, 768], row t = wte[inputIds[t]]
}

Don't transpose or normalize wte

wte is already stored [vocab, 768] row-major — row id is the embedding. Index it directly; do not transpose it.

And do not scale. The original Transformer multiplied embeddings by √d_model; GPT-2 does not. Multiply by √768 ≈ 27.7 and every value blows up — your output will miss the reference by a mile.

How it's graded

Your [T, 768] output is compared to the reference with allclose, atol 1e-5.

Ids to vectors Add positions