NNX Implement Positional Embed

Give each position a vector

A token embedding describes which token is present. A learned positional embedding supplies a separate vector for its location: row 0 for the first token, row 1 for the second, and so on. Adding the two gives the next layer information about both token identity and position.

Build a trainable table with shape (64, d_model). For a sequence of length T, return its first T rows. This exercise covers learned absolute positions starting at zero. Other designs use explicit position IDs, offsets during generation, or rotary position encodings.

Initialization and forward pass

Store the table in an nnx.Param so NNX's parameter filters include it. Draw one key from rngs.params(), sample a float32 standard-normal array of shape (64, d_model), and multiply by 0.02. This scale is part of the exercise's initialization recipe. The relative size of token and position signals also depends on how the token embeddings are initialized.

The forward pass only slices the table. It does not draw another random key or create new parameters. During training, a loss using this slice can produce gradients for rows 0 through T - 1; unused rows receive zero gradient from this lookup. An optimizer with weight decay or stored momentum can still change parameters with zero current gradient.

Allocate all 64 rows even for a short sequence. Changing the random sample shape to (T, d_model) changes the specified initialization. With the same seed and width, calls to the entrypoint with different valid lengths must return prefixes of the same initialized table.

Worked example: adding positions

This complete example replaces the random values with a small table so the addition is easy to check. Repeated token vectors become different at different positions.

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

class MyPositionalEmbed(nnx.Module):
    def __init__(self, max_T, d_model, rngs):
        key = rngs.params()
        self.pos_embed = nnx.Param(
            jax.random.normal(key, (max_T, d_model), dtype=jnp.float32)
            * 0.02
        )

    def __call__(self, T):
        return self.pos_embed[:T]

model = MyPositionalEmbed(3, 2, nnx.Rngs(0))
model.pos_embed[...] = jnp.array(
    [[0.1, 0.0], [0.0, 0.2], [-0.1, 0.3]], dtype=jnp.float32
)
tokens = jnp.array([[1.0, 2.0], [1.0, 2.0]], dtype=jnp.float32)
combined = tokens + model(2)
print(combined)  # [[1.1, 2.0], [1.0, 2.2]]

For a batch of token vectors shaped (batch, T, d_model), the same (T, d_model) position slice broadcasts across the batch. The exercise itself returns only the position vectors, without adding token vectors.

Sequence-length limits

Valid inputs satisfy 1 <= T <= 64. A slice such as table[:65] on a 64-row table returns 64 rows; it does not raise an error or invent a new position vector. The slice is therefore not input validation. A production model should explicitly handle lengths beyond its supported range.

A larger table also does not mean every position has been trained. Rows that training examples never select receive no learning signal from those lookups.

Compare with fixed positions

The visualization below shows sinusoidal position encodings. Those values come from a fixed formula. They illustrate position-dependent vectors, but are not the learned random table you implement here.