Foundations

Token Embedding Lookup

EasyRaw PyTorchAbout 10 min1 public example plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Look up each integer token in the supplied embedding table. Return one embedding vector per token, optionally zeroing the configured padding token.

Function signature

embedding_lookup(
    tokens: torch.Tensor,
    table: torch.Tensor,
    padding_id: None | int = None,
) -> torch.Tensor

Worked example

Repeated tokens and padding

Token 2 selects table row [-2, 4] twice; token 1 selects [1, 3]. Token 0 is configured as padding, so its output is [0, 0] instead of table row [9, 8].

import torch

tokens = torch.tensor([[2, 1, 2, 0]], dtype=torch.int64)
table = torch.tensor([[9.0, 8.0], [1.0, 3.0], [-2.0, 4.0]], dtype=torch.float64, requires_grad=True)
padding_id = 0

result = embedding_lookup(tokens=tokens, table=table, padding_id=padding_id)

Expected return

torch.tensor([[[-2.0, 4.0], [1.0, 3.0], [-2.0, 4.0], [0.0, 0.0]]], dtype=torch.float64)

Parameters

Parameter	Shape and meaning
`tokens`torch.Tensor	$[\ldots]$ Integer token IDs indexing rows of the embedding table.
`table`torch.Tensor	$[V, C]$ Supplied vocabulary embedding table; no new parameters are initialized.
`padding_id`None \| intDefault `None`	Optional token ID to suppress. None keeps every looked-up vector; otherwise matching token positions return zero.

Return

torch.Tensor · $[\ldots,C]$

Return One embedding vector for every token position.

Required behavior

tokens is an integer tensor of any leading shape with IDs in[0,V); table has shape $[V,C]$ . Return shape (*tokens.shape,C).
padding_id is optional; None disables padding masking. When set, entries with that token ID are exactly zero.
Do not create a new table or draw random values. Repeated IDs reuse the same supplied row and their gradients accumulate into it.
The return dtype/device is the table dtype/device. Token IDs are nondifferentiable.
Use CPU float32/float64 tensors with one shared floating dtype/device. Accept noncontiguous input layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

Why this matters

A projection mixes a vector of features; a lookup selects a learned vector using an integer identity.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def embedding_lookup(
    tokens: torch.Tensor,  # (...)
    table: torch.Tensor,  # (V, C)
    padding_id: None | int = None,
) -> torch.Tensor:
    """Return a tensor with shape (...,C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement embedding_lookup")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.