Foundations · INTEGRATION

Affine projection and embedding lookup

EasyRaw PyTorchAbout 10 min2 public examples plus private tests

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

Your task

Implement either a final-channel affine projection or an integer embedding lookup. Preserve every leading input axis and return the projected or gathered feature tensor.

Function signature

affine_projection(
    x: torch.Tensor | None = None,
    weight: torch.Tensor | None = None,
    bias: torch.Tensor | None = None,
    tokens: torch.Tensor | None = None,
    table: torch.Tensor | None = None,
    variant: str = 'affine',
    padding_id: None | int = None,
) -> torch.Tensor

Worked example

One projected feature vector

For variant="affine", x=[1,2], weight=[[1,3],[2,-1]] and bias=[0.5,1], the two output channels are 1×1+2×2+0.5=5.5 and 1×3+2×(-1)+1=2. Return [5.5,2]. The weight columns determine the output channels; the embedding variant below instead selects table rows.

Parameters

Parameter	Shape and meaning
`x`torch.Tensor \| NoneDefault `None`	$[\ldots, \mathrm{Cin}]$ Input feature values to transform.
`weight`torch.Tensor \| NoneDefault `None`	$[\mathrm{Cin}, \mathrm{Cout}]$ Affine only: floating $[\mathrm{Cin},\mathrm{Cout}]$ , not PyTorch Linear's stored orientation.
`bias`torch.Tensor \| NoneDefault `None`	$[\mathrm{Cout}]$ Optional affine `bias` $[\mathrm{Cout}]$ .
`tokens`torch.Tensor \| NoneDefault `None`	$[\ldots]$ Embedding only: integer $[\ldots]$ in $[0,V)$ .
`table`torch.Tensor \| NoneDefault `None`	$[V, C]$ Embedding only: floating learned $[V,C]$ .
`variant`strDefault `'affine'`	Choose "affine"/"linear" for channel projection or "embedding"/"lookup" for integer table lookup.
`padding_id`None \| intDefault `None`	Embedding mode only: an optional integer token ID whose gathered output is zero. None disables padding suppression.

Return

torch.Tensor · Affine: $[\ldots, \mathrm{Cout}]$ . Embedding: $[\ldots, C]$ . The leading axes come from x or tokens, respectively.

Return Affine [...,Cout] or embedding [...,C].

Required behavior

Implement one operation selected by variant (default 'affine'; 'linear' is an alias).
For affine/linear, $x$ has shape $[\ldots,\mathrm{Cin}]$ , weight has shape $[\mathrm{Cin},\mathrm{Cout}]$ , and optional bias has shape $[\mathrm{Cout}]$ . Return $\mathrm{output}[\ldots,o] = \sum_{i} x[\ldots,i]\cdot \mathrm{weight}[i,o] + \mathrm{bias}[o]$ , omitting bias when absent; the final feature axis alone is contracted.
For variant 'embedding' (alias 'lookup'), integer tokens $[\ldots]$ lie in $[0,V)$ and table $[V,C]$ is supplied; return $\mathrm{output}[\ldots,c] = \mathrm{table}[\mathrm{tokens}[\ldots],c]$ .
Optional padding_id is either None (default, all IDs ordinary) or one ID in $[0,V)$ ; its output vector is forced to zero even if the supplied table row is nonzero, and that row receives zero gradient from padded positions. Duplicate nonpadding IDs accumulate table gradients.
Preserve dtype/device and arbitrary leading axes; accept noncontiguous inputs; empty token axes are allowed. Preserve the declared return shape; do not mutate inputs, initialize weights, normalize embeddings or add residuals.

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 affine_projection(
    x: torch.Tensor | None = None,  # (..., Cin)
    weight: torch.Tensor | None = None,  # (Cin, Cout)
    bias: torch.Tensor | None = None,  # (Cout,)
    tokens: torch.Tensor | None = None,  # (...)
    table: torch.Tensor | None = None,  # (V, C)
    variant: str = 'affine',
    padding_id: None | int = None,
) -> torch.Tensor:
    """Return a tensor with shape Affine: (..., Cout). Embedding: (..., C). The leading
        axes come from x or tokens, respectively."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement affine_projection")

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.