numpy_core.py

  1"""Transparent forward passes for the Latent Space attention course.
  2
  3Install numpy. Run: python numpy_core.py
  4Regions are included directly in the lessons, so displayed code stays in sync.
  5"""
  6from __future__ import annotations
  7
  8import numpy as np
  9from numpy.typing import NDArray
 10
 11FloatArray = NDArray[np.float64]
 12IntArray = NDArray[np.int64]
 13BoolArray = NDArray[np.bool_]
 14# Type aliases describe dtype; docstrings and assertions specify tensor shapes.
 15
 16# region embeddings
 17def embedding_demo() -> tuple[IntArray, FloatArray, FloatArray]:
 18    vocab = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "the": 3,
 19             "small": 4, "robot": 5, "moves": 6, "<unk>": 7}
 20    tokens = "the small robot moves".split()  # Teaching tokenizer only.
 21    ids = np.array([vocab.get(token, vocab["<unk>"]) for token in tokens])
 22    table = np.random.default_rng(0).normal(size=(len(vocab), 6))
 23    x = table[ids]                           # (4, 6), no arithmetic on IDs.
 24    assert np.array_equal(x[2], table[vocab["robot"]])
 25    return ids, table, x
 26# endregion embeddings
 27
 28# region positions
 29def positions(length: int, width: int, start: int = 0) -> FloatArray:
 30    if width <= 0 or width % 2:
 31        raise ValueError("This example uses a positive, even model width")
 32    pos = np.arange(start, start + length, dtype=float)[:, None]
 33    frequency = 10000.0 ** (-np.arange(0, width, 2) / width)
 34    angles = pos * frequency[None, :]
 35    pe = np.empty((length, width))
 36    pe[:, 0::2] = np.sin(angles)
 37    pe[:, 1::2] = np.cos(angles)
 38    return pe
 39# endregion positions
 40
 41# region attention
 42def attention(q: FloatArray, k: FloatArray, v: FloatArray, allowed: BoolArray | None = None) -> tuple[FloatArray, FloatArray]:
 43    """(..., T, dk), (..., S, dk), (..., S, dv). True means visible."""
 44    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
 45        raise ValueError("Q/K feature widths and K/V source lengths must match")
 46    scores = (q @ np.swapaxes(k, -2, -1)) / np.sqrt(q.shape[-1])
 47    if allowed is not None:
 48        visible = np.broadcast_to(np.asarray(allowed, dtype=bool), scores.shape)
 49        if not visible.any(axis=-1).all():
 50            raise ValueError("Every query needs at least one visible key")
 51        scores = np.where(visible, scores, -np.inf)
 52    shifted = scores - scores.max(axis=-1, keepdims=True)
 53    weights = np.exp(shifted)
 54    weights /= weights.sum(axis=-1, keepdims=True)
 55    return weights @ v, weights
 56# endregion attention
 57
 58# region masks
 59def attention_masks(source_ids: IntArray, target_ids: IntArray, pad: int = 0) -> tuple[BoolArray, BoolArray]:
 60    """IDs: (batch, length). Return source-key and causal-target masks."""
 61    source_visible = (source_ids != pad)[:, None, None, :]
 62    target_visible = (target_ids != pad)[:, None, None, :]
 63    length = target_ids.shape[1]
 64    causal = np.arange(length)[None, :] <= np.arange(length)[:, None]
 65    target_visible = target_visible & causal[None, None, :, :]
 66    return source_visible, target_visible
 67# endregion masks
 68
 69# region heads
 70class MultiHead:
 71    def __init__(self, width: int, heads: int, seed: int = 0) -> None:
 72        if heads <= 0 or width % heads:
 73            raise ValueError("Model width must be divisible by head count")
 74        self.heads, self.head_width = heads, width // heads
 75        rng = np.random.default_rng(seed)
 76        self.wq, self.wk, self.wv, self.wo = [
 77            rng.normal(size=(width, width)) / np.sqrt(width) for _ in range(4)
 78        ]
 79
 80    def split(self, x: FloatArray) -> FloatArray:
 81        batch, length, _ = x.shape
 82        return x.reshape(batch, length, self.heads, self.head_width).transpose(0, 2, 1, 3)
 83
 84    def __call__(self, query: FloatArray, context: FloatArray, allowed: BoolArray | None = None) -> tuple[FloatArray, FloatArray]:
 85        q = self.split(query @ self.wq)
 86        k = self.split(context @ self.wk)
 87        v = self.split(context @ self.wv)
 88        heads, weights = attention(q, k, v, allowed)
 89        batch, _, length, _ = heads.shape
 90        joined = heads.transpose(0, 2, 1, 3).reshape(batch, length, -1)
 91        return joined @ self.wo, weights
 92# endregion heads
 93
 94# region norm_ffn
 95def layer_norm(x: FloatArray, gamma: FloatArray | None = None, beta: FloatArray | None = None, eps: float = 1e-5) -> FloatArray:
 96    mean = x.mean(axis=-1, keepdims=True)
 97    variance = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
 98    normalized = (x - mean) / np.sqrt(variance + eps)
 99    gamma = np.ones(x.shape[-1]) if gamma is None else gamma
100    beta = np.zeros(x.shape[-1]) if beta is None else beta
101    return normalized * gamma + beta
102
103
104def feed_forward(x: FloatArray, w1: FloatArray, b1: FloatArray, w2: FloatArray, b2: FloatArray) -> FloatArray:
105    hidden = np.maximum(x @ w1 + b1, 0.0)  # ReLU, separately per token.
106    return hidden @ w2 + b2               # Return to the residual width.
107# endregion norm_ffn
108
109
110def ffn_parameters(width: int, inner: int, seed: int) -> tuple[FloatArray, FloatArray, FloatArray, FloatArray]:
111    rng = np.random.default_rng(seed)
112    return (rng.normal(size=(width, inner)) / np.sqrt(width), np.zeros(inner),
113            rng.normal(size=(inner, width)) / np.sqrt(inner), np.zeros(width))
114
115# region encoder
116class EncoderBlock:
117    """A pre-norm forward pass. Normalization gain=1 and bias=0 here."""
118    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 10) -> None:
119        self.attn = MultiHead(width, heads, seed)
120        self.ffn = ffn_parameters(width, inner, seed + 1)
121
122    def __call__(self, x: FloatArray, source_visible: BoolArray) -> FloatArray:
123        normalized = layer_norm(x)
124        x = x + self.attn(normalized, normalized, source_visible)[0]
125        return x + feed_forward(layer_norm(x), *self.ffn)
126# endregion encoder
127
128# region decoder
129class DecoderBlock:
130    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 20) -> None:
131        self.self_attn = MultiHead(width, heads, seed)
132        self.cross_attn = MultiHead(width, heads, seed + 1)
133        self.ffn = ffn_parameters(width, inner, seed + 2)
134
135    def __call__(self, y: FloatArray, memory: FloatArray, target_visible: BoolArray, source_visible: BoolArray) -> FloatArray:
136        normalized = layer_norm(y)
137        y = y + self.self_attn(normalized, normalized, target_visible)[0]
138        y = y + self.cross_attn(layer_norm(y), memory, source_visible)[0]
139        return y + feed_forward(layer_norm(y), *self.ffn)
140# endregion decoder
141
142# region architectures
143def architecture_demo() -> tuple[FloatArray, FloatArray]:
144    source = np.array([[4, 7, 9], [6, 8, 0]])
145    prefix = np.array([[1, 4, 7], [1, 6, 8]])
146    width, vocabulary = 32, 12
147    rng = np.random.default_rng(8)
148    embedding = rng.normal(size=(vocabulary, width))
149    source_mask, target_mask = attention_masks(source, prefix)
150    x = embedding[source] + positions(source.shape[1], width)
151    y = embedding[prefix] + positions(prefix.shape[1], width)
152
153    # Encoder-only: pool only real source positions for classification.
154    encoded = layer_norm(EncoderBlock()(x, source_mask))
155    real = (source != 0)[..., None]
156    pooled = (encoded * real).sum(1) / real.sum(1)
157    class_logits = pooled @ rng.normal(size=(width, 3))
158
159    # Decoder-only: same two-sublayer structure, with CAUSAL visibility.
160    causal_hidden = layer_norm(EncoderBlock(seed=30)(y, target_mask))
161    next_token_logits = causal_hidden @ rng.normal(size=(width, vocabulary))
162    return class_logits, next_token_logits  # (2, 3), (2, 3, 12)
163# endregion architectures
164
165# region forward
166def demo_forward() -> FloatArray:
167    # PAD=0, BOS=1, EOS=2; source symbols are integers 4..11.
168    source = np.array([[4, 7, 9], [6, 8, 0]])
169    target = np.array([[1, 4, 7], [1, 6, 8]])
170    width, vocabulary = 32, 12
171    rng = np.random.default_rng(7)
172    embedding = rng.normal(size=(vocabulary, width))
173    source_mask, target_mask = attention_masks(source, target)
174    x = embedding[source] + positions(source.shape[1], width)
175    y = embedding[target] + positions(target.shape[1], width)
176    memory = layer_norm(EncoderBlock()(x, source_mask))
177    hidden = layer_norm(DecoderBlock()(y, memory, target_mask, source_mask))
178    output_projection = rng.normal(size=(width, vocabulary)) / np.sqrt(width)
179    logits = hidden @ output_projection
180    assert logits.shape == (2, 3, vocabulary)
181    return logits  # Random, untrained scores: not meaningful predictions.
182# endregion forward
183
184if __name__ == "__main__":
185    print("Embedding shape:", embedding_demo()[2].shape)
186    print("Vocabulary logits:", demo_forward().shape)