foundations_numpy.py

 1"""Small typed experiments for the foundations course. Run this file directly."""
 2from __future__ import annotations
 3
 4import numpy as np
 5from numpy.typing import NDArray
 6
 7FloatArray = NDArray[np.float64]
 8
 9# region shapes
10def project(x: FloatArray, weight: FloatArray, bias: FloatArray) -> FloatArray:
11    """x: (tokens, input_features), weight: (input_features, output_features)."""
12    if x.ndim != 2 or weight.ndim != 2 or x.shape[1] != weight.shape[0]:
13        raise ValueError("The input feature axis must match the weight's row axis")
14    products = x @ weight
15    output = products + bias  # (output_features,) broadcasts over tokens.
16    return output
17# endregion shapes
18
19# region mixture
20def weighted_mix(weights: FloatArray, values: FloatArray) -> FloatArray:
21    """weights: (sources,), values: (sources, features); result: (features,)."""
22    if np.any(weights < 0) or not np.isclose(weights.sum(), 1.0):
23        raise ValueError("Use nonnegative weights summing to one")
24    contributions = weights[:, None] * values
25    output = contributions.sum(axis=0)
26    return output
27# endregion mixture
28
29# region probability
30def softmax(scores: FloatArray, temperature: float = 1.0) -> FloatArray:
31    """Normalize the final axis; temperature must be positive."""
32    if temperature <= 0:
33        raise ValueError("Temperature must be positive")
34    scaled = scores / temperature
35    shifted = scaled - scaled.max(axis=-1, keepdims=True)
36    evidence = np.exp(shifted)
37    probabilities = evidence / evidence.sum(axis=-1, keepdims=True)
38    return probabilities
39# endregion probability
40
41# region learning
42def learn_step(weight: float, x: float, target: float, rate: float = 0.1) -> tuple[float, float, float]:
43    """One squared-error gradient step for prediction = weight * x."""
44    prediction = weight * x
45    error = prediction - target
46    loss = error ** 2
47    gradient = 2.0 * error * x
48    updated = weight - rate * gradient
49    return updated, loss, gradient
50# endregion learning
51
52
53def check() -> None:
54    x = np.array([[1.0, 2.0], [3.0, 4.0]])
55    weight = np.array([[1.0, 0.0], [0.5, 2.0]])
56    np.testing.assert_allclose(project(x, weight, np.array([0.0, 1.0])), [[2, 5], [5, 9]])
57    np.testing.assert_allclose(weighted_mix(np.array([0.2, 0.3, 0.5]), np.array([[1., 0.], [0., 2.], [2., 1.]])), [1.2, 1.1])
58    np.testing.assert_allclose(softmax(np.array([0., np.log(2.), np.log(3.)])), [1/6, 2/6, 3/6])
59    np.testing.assert_allclose(learn_step(0.5, 2.0, 3.0), [1.3, 4.0, -8.0])
60    print("PASS: foundations projections, mixtures, softmax, and gradient update")
61
62
63if __name__ == "__main__":
64    check()