NNX Implement Embed

Why this matters

An embedding maps an integer ID to a learned vector. Store the vectors as rows of a table with shape (vocab_size, d_model): the first dimension counts possible IDs and the second gives the vector length. Looking up T IDs produces an array of shape (T, d_model).

The table contains vocab_size * d_model scalar parameters. Its share of an entire model's parameters depends on the vocabulary, width and other layers; an embedding is not necessarily the largest part of a model.

A complete example

import jax
import jax.numpy as jnp
from flax import nnx

class MyEmbed(nnx.Module):
    def __init__(self, vocab_size, d_model, rngs):
        key = rngs.params()
        self.embedding = nnx.Param(
            jax.random.normal(key, (vocab_size, d_model), dtype=jnp.float32)
        )

    def __call__(self, token_ids):
        return self.embedding[token_ids]

model = MyEmbed(3, 2, nnx.Rngs(0))
assert model.embedding[...].shape == (3, 2)

# Known rows make the lookup order easy to see.
model.embedding[...] = jnp.array([[1., 2.], [3., 4.], [5., 6.]])
ids = jnp.array([2, 0, 2], dtype=jnp.int32)
expected = jnp.array([[5., 6.], [1., 2.], [5., 6.]])
assert jnp.array_equal(model(ids), expected)
assert jnp.array_equal(model(ids).reshape(-1), jnp.array([5., 6., 1., 2., 5., 6.]))

The repeated ID selects the same row twice, in the requested order. The final assignment to known rows is only for demonstration; use the random table when solving the exercise.

Initialization and learning

Initialize the table with one float32 standard-normal draw from rngs.params(), without an extra scale factor. This is the exercise's initialization contract, not a recommendation for every model. Different scales or built-in embedding defaults can produce different seeded values.

Indexing selects existing parameters; it does not generate a new vector per token occurrence. During differentiation, contributions from repeated IDs accumulate into the same table row. Gradients from the lookup itself reach selected rows; unselected rows receive zero. This does not mean an optimizer with weight decay or stored momentum can never change those rows. Integer IDs are selections, not continuous differentiable inputs.

IDs must be valid integers

The harness represents IDs as float32 values, so cast them to jnp.int32 before indexing. The input contract guarantees they are integer-valued and lie between 0 and vocab_size - 1. Casting is not validation: it can truncate a fractional value, and indexing does not reliably reject every invalid ID with an exception. Do not use out-of-range behavior as part of the layer's semantics.

Direct row gathering avoids constructing a separate one-hot table for the IDs. A one-hot matrix product describes the same lookup algebra for valid IDs, but this exercise asks you to use indexing.