The projection head·25 min

The full encode

0 / 1

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

Now compose the whole projection head for a batch of already-pooled sentences: linear1 → tanh → linear2 → L2-normalize. The output is the [B, 32] matrix of unit-length sentence embeddings — the thing the contrastive loss will operate on.

one row of the batch

pooled[b]  ──linear(W1,b1)──▶  h1  ──tanh──▶  a1
           ──linear(W2,b2)──▶  h2  ──L2-norm──▶  out[b]   (‖out[b]‖ = 1)

What to build

Your task

Implement encode

Implement runEncode(pooled, W1, b1, W2, b2): given a [B, 32] pooled batch and the MLP weights, return the [B, 32] matrix of L2-normalized sentence embeddings. Apply the head row by row, reusing your linear and tanh.

tests/adapters.ts

export function runEncode(
  pooled: number[][], W1: number[][], b1: number[], W2: number[][], b2: number[],
): number[][] {
  // each row: linear(W1,b1) -> tanh -> linear(W2,b2) -> l2normalize
}

How it's graded

Adapter mode: your [B, 32] matrix is compared to the reference with allclose, atol 1e-5, rtol 1e-4. The fixture uses the frozen init weights and a fixed pooled batch.

Order matters: normalize last

L2-normalize is the final step, after the second linear. Normalizing earlier (e.g. right after pooling, or between the linears) changes the math and the gradient — the reference normalizes exactly once, at the very end.

Linear + tanh Recap: the encoder