| """V20-sized learned controller for the two-program U2 arithmetic board.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from src.self_clocked_parallel_montgomery_teacher_student import select_candidate |
| from src.universal_self_clocked_u1 import OP_COUNT, PHASE_COUNT, UniversalU1Memory |
| from src.universal_self_clocked_u2 import ( |
| ALGORITHM_COUNT, |
| UniversalSelfClockedU2, |
| ) |
|
|
|
|
| class CompactMixedAlgorithmSelfClockTeacher(nn.Module): |
| """Learn a finite program using compact scalar clock transitions.""" |
|
|
| def __init__(self, *, maximum_width: int = 512, maximum_clock: int = 2048) -> None: |
| super().__init__() |
| if maximum_width < 1: |
| raise ValueError("maximum_width must be positive") |
| if maximum_clock < 4 * maximum_width: |
| raise ValueError("maximum_clock must cover four bits per base-16 digit") |
| self.maximum_width = int(maximum_width) |
| self.maximum_clock = int(maximum_clock) |
| self.operation = nn.Embedding(PHASE_COUNT * 2, OP_COUNT) |
| self.initial_phase = nn.Embedding(ALGORITHM_COUNT, PHASE_COUNT) |
| self.initial_clock_affine = nn.Embedding(ALGORITHM_COUNT, 2) |
| self.decrement_clock_affine = nn.Linear(1, 1) |
| self.terminal_clock = nn.Parameter(torch.empty(1)) |
| self.terminal_test = nn.Linear(1, 2) |
| self.next_phase = nn.Embedding(PHASE_COUNT * OP_COUNT, PHASE_COUNT) |
| for parameter in self.parameters(): |
| nn.init.normal_(parameter, std=0.02) |
| self.config = { |
| "maximum_width": self.maximum_width, |
| "maximum_clock": self.maximum_clock, |
| } |
|
|
| @property |
| def parameter_count(self) -> int: |
| return sum(parameter.numel() for parameter in self.parameters()) |
|
|
| def _validate_algorithms(self, algorithm_ids: torch.Tensor) -> None: |
| if algorithm_ids.ndim != 1: |
| raise ValueError("algorithm_ids must be a one-dimensional batch") |
| if bool( |
| ((algorithm_ids < 0) | (algorithm_ids >= ALGORITHM_COUNT)).any().item() |
| ): |
| raise ValueError("algorithm id is outside the learned instruction domain") |
|
|
| def initial_clock_values( |
| self, algorithm_ids: torch.Tensor, widths: torch.Tensor |
| ) -> torch.Tensor: |
| self._validate_algorithms(algorithm_ids) |
| if algorithm_ids.shape != widths.shape: |
| raise ValueError("algorithm ids and widths must align") |
| scale_bias = self.initial_clock_affine(algorithm_ids) |
| return scale_bias[:, 0] * widths.to(scale_bias.dtype) + scale_bias[:, 1] |
|
|
| def raw_decrement_clock_values(self, clocks: torch.Tensor) -> torch.Tensor: |
| return self.decrement_clock_affine( |
| clocks.to(torch.float32).unsqueeze(-1) |
| ).squeeze(-1) |
|
|
| def decrement_clock_values(self, clocks: torch.Tensor) -> torch.Tensor: |
| return F.relu(self.raw_decrement_clock_values(clocks)) |
|
|
| def terminal_test_logits(self, clocks: torch.Tensor) -> torch.Tensor: |
| return self.terminal_test(clocks.to(torch.float32).unsqueeze(-1)) |
|
|
| def _discrete_clock(self, values: torch.Tensor) -> torch.Tensor: |
| return values.round().clamp(0, self.maximum_clock).to(torch.long) |
|
|
| def start( |
| self, width: int, algorithm_ids: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| if not 1 <= int(width) <= self.maximum_width: |
| raise ValueError("width exceeds the learned controller domain") |
| self._validate_algorithms(algorithm_ids) |
| widths = torch.full_like(algorithm_ids, int(width)) |
| phase = self.initial_phase(algorithm_ids).argmax(-1) |
| clock = self._discrete_clock( |
| self.initial_clock_values(algorithm_ids, widths) |
| ) |
| return phase, clock |
|
|
| def operation_logits( |
| self, phases: torch.Tensor, clocks: torch.Tensor |
| ) -> torch.Tensor: |
| clock_class = self.terminal_test_logits(clocks).argmax(-1) |
| return self.operation(phases * 2 + clock_class) |
|
|
| def instruction(self, phase: torch.Tensor, clock: torch.Tensor) -> torch.Tensor: |
| return self.operation_logits(phase, clock).argmax(-1) |
|
|
| def advance( |
| self, |
| phase: torch.Tensor, |
| clock: torch.Tensor, |
| operation: torch.Tensor, |
| algorithm_ids: torch.Tensor, |
| *, |
| width: int, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| widths = torch.full_like(clock, int(width)) |
| row_clock = self._discrete_clock(self.decrement_clock_values(clock)) |
| reset_clock = self._discrete_clock( |
| self.initial_clock_values(algorithm_ids, widths) |
| ) |
| terminal_clock = self._discrete_clock(self.terminal_clock).expand_as(clock) |
| next_clock = select_candidate( |
| ( |
| row_clock, |
| reset_clock, |
| terminal_clock, |
| row_clock, |
| terminal_clock, |
| clock, |
| ), |
| operation, |
| ) |
| next_phase = self.next_phase(phase * OP_COUNT + operation).argmax(-1) |
| return next_phase, next_clock |
|
|
|
|
| class V20CompactUniversalSelfClockedU2(UniversalSelfClockedU2): |
| """Use the U2 arithmetic board with a compact learned controller.""" |
|
|
| def __init__( |
| self, |
| *, |
| reducer: dict, |
| controller: dict, |
| mac: dict, |
| subtract: dict, |
| factor: dict, |
| cancellation_teacher: dict, |
| prime_two_router: dict, |
| ) -> None: |
| super().__init__( |
| reducer=reducer, |
| controller=controller, |
| mac=mac, |
| subtract=subtract, |
| factor=factor, |
| cancellation_teacher=cancellation_teacher, |
| prime_two_router=prime_two_router, |
| ) |
| self.controller = CompactMixedAlgorithmSelfClockTeacher(**controller) |
| self.config = { |
| "reducer": dict(reducer), |
| "controller": dict(controller), |
| "mac": dict(mac), |
| "subtract": dict(subtract), |
| "factor": dict(factor), |
| "cancellation_teacher": dict(cancellation_teacher), |
| "prime_two_router": dict(prime_two_router), |
| } |
|
|
| def _controller_instruction(self, memory: UniversalU1Memory) -> torch.Tensor: |
| return self.controller.instruction(memory.phase, memory.clock) |
|
|
|
|
| __all__ = [ |
| "CompactMixedAlgorithmSelfClockTeacher", |
| "V20CompactUniversalSelfClockedU2", |
| ] |
|
|
