foundations_torch.py

 1"""Tensor versions of the foundations experiments, with automatic gradients."""
 2from __future__ import annotations
 3
 4import torch
 5from torch import Tensor
 6
 7# region shapes
 8def project(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
 9    """x: (tokens, input_features), weight: (input_features, output_features)."""
10    if x.ndim != 2 or weight.ndim != 2 or x.shape[1] != weight.shape[0]:
11        raise ValueError("The input feature axis must match the weight's row axis")
12    products = x @ weight
13    output = products + bias
14    return output
15# endregion shapes
16
17# region mixture
18def weighted_mix(weights: Tensor, values: Tensor) -> Tensor:
19    """weights: (sources,), values: (sources, features); result: (features,)."""
20    if torch.any(weights < 0) or not torch.isclose(weights.sum(), weights.new_tensor(1.0)):
21        raise ValueError("Use nonnegative weights summing to one")
22    contributions = weights[:, None] * values
23    output = contributions.sum(dim=0)
24    return output
25# endregion mixture
26
27# region probability
28def softmax(scores: Tensor, temperature: float = 1.0) -> Tensor:
29    if temperature <= 0:
30        raise ValueError("Temperature must be positive")
31    scaled = scores / temperature
32    probabilities = scaled.softmax(dim=-1)
33    return probabilities
34# endregion probability
35
36# region learning
37def learn_step(weight: float, x: float, target: float, rate: float = 0.1) -> tuple[float, float, float]:
38    parameter = torch.tensor(weight, dtype=torch.float64, requires_grad=True)
39    prediction = parameter * x
40    loss = (prediction - target) ** 2
41    loss.backward()
42    assert parameter.grad is not None
43    gradient = parameter.grad.item()
44    updated = parameter.item() - rate * gradient
45    return updated, loss.item(), gradient
46# endregion learning
47
48
49def check() -> None:
50    import foundations_numpy as reference
51    import numpy as np
52
53    x = np.array([[1., 2.], [3., 4.]])
54    w = np.array([[1., 0.], [0.5, 2.]])
55    b = np.array([0., 1.])
56    np.testing.assert_allclose(project(torch.tensor(x), torch.tensor(w), torch.tensor(b)).numpy(), reference.project(x, w, b))
57    np.testing.assert_allclose(learn_step(0.5, 2.0, 3.0), reference.learn_step(0.5, 2.0, 3.0))
58    print("PASS: foundations NumPy/PyTorch agreement")
59
60
61if __name__ == "__main__":
62    check()