Skip to content

← Flax step 31 of 100

Medium Primitives

Token Embedding with Flax

Jump to the task

From token IDs to vectors

A token ID identifies an item in a vocabulary. An embedding layer gives each ID a vector that later layers can process. Its table W has shape (V, D): V vocabulary entries, each with D features.

The lookup is out[t] = W[ids[t]]. A sequence of T IDs therefore produces shape (T, D). Repeated IDs select the same row, and changing the order of IDs changes the output order. The lookup itself adds no information about a token's position or neighboring tokens.

This exercise initializes a table and reads from it. It does not train the vectors or load pretrained embeddings.

A lookup you can check by hand

A one-hot vector has a 1 at the selected ID and zeros elsewhere. Multiplying it by W selects that row. This example supplies known weights so the result is easy to inspect:

import jax
import jax.numpy as jnp
import flax.linen as nn

table = jnp.array([[1., 2.], [3., 4.], [5., 6.]])
ids = jnp.array([2, 0, 2], dtype=jnp.int32)
embed = nn.Embed(num_embeddings=3, features=2)
variables = {"params": {"embedding": table}}
out = embed.apply(variables, ids)

assert jnp.array_equal(out, jnp.array([[5., 6.], [1., 2.], [5., 6.]]))
assert jnp.array_equal(out, jax.nn.one_hot(ids, 3) @ table)

Direct lookup avoids constructing the (T, V) one-hot array. It still stores the full (V, D) embedding table.

Initialize and apply a Linen module

Constructing nn.Embed describes the layer. Call init to create its variables, then pass those variables to apply:

import jax
import jax.numpy as jnp
import flax.linen as nn

ids = jnp.array([0, 2, 1], dtype=jnp.int32)
embed = nn.Embed(num_embeddings=4, features=3)
variables = embed.init(jax.random.PRNGKey(0), ids)
table = variables["params"]["embedding"]
out = embed.apply(variables, ids)

assert table.shape == (4, 3)
assert out.shape == (3, 3)
assert jnp.array_equal(out, table[ids])

Keep Linen's default initializer for this exercise. It samples a normal distribution with variance 1 / D, so its standard deviation is 1 / sqrt(D). Linen derives the parameter's random key internally; drawing directly from the input key does not reproduce init.

What learning would change

If a loss depends on the table only through this lookup, its gradient is zero for rows that were not selected. Contributions from repeated IDs add together. This does not guarantee that an optimizer leaves unused rows unchanged: weight decay or stored momentum can update them. Sharing the table with an output classifier can also give those rows gradients. No optimizer runs in this exercise.

Problem

Implement token_embed_forward(seed, token_ids, vocab_size, d_model):

  1. Convert seed, vocab_size and d_model to Python integers.
  2. Cast token_ids to jnp.int32.
  3. Construct nn.Embed(num_embeddings=vocab_size, features=d_model) with its default initializer.
  4. Initialize it with jax.random.PRNGKey(seed) and the integer IDs.
  5. Apply it with the returned variables and the same IDs.
  6. Flatten the (T, D) result with .reshape(-1) and return it.

vocab_size and d_model are positive. The harness supplies token_ids as a one-dimensional float array of whole numbers in [0, vocab_size). Linen requires an integer dtype, which is why the cast is needed. Casting is not validation for fractional or out-of-range IDs; those inputs are outside this problem's contract.

The output has length T * D, with each token's features kept together. If you wrap the function in jax.jit, make seed, vocab_size and d_model static arguments; the ID array can remain dynamic.