| """Learned raw-input algorithm selection for the exact U2 executor. |
| |
| This research wrapper contains one learned router and one frozen U2 executor. |
| The router sees only raw base-16 digits, argument roles, and lengths. It emits |
| the U2 algorithm instruction; it does not compute or inspect the answer. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from src.universal_self_clocked_u1 import UniversalU1Prediction |
| from src.universal_self_clocked_u2 import ( |
| ALGORITHM_COUNT, |
| UniversalSelfClockedU2, |
| ) |
|
|
|
|
| class RawInputAlgorithmRouter(nn.Module): |
| """Small input-capable router over three raw base-16 rows.""" |
|
|
| def __init__( |
| self, |
| *, |
| maximum_digits: int = 1024, |
| digit_dimension: int = 16, |
| length_dimension: int = 8, |
| hidden_dimension: int = 64, |
| algorithm_count: int = ALGORITHM_COUNT, |
| feature_mode: str = "token_summary", |
| ) -> None: |
| super().__init__() |
| if min( |
| maximum_digits, |
| digit_dimension, |
| length_dimension, |
| hidden_dimension, |
| algorithm_count, |
| ) < 1: |
| raise ValueError("all router dimensions must be positive") |
| self.maximum_digits = int(maximum_digits) |
| self.digit_dimension = int(digit_dimension) |
| self.length_dimension = int(length_dimension) |
| self.hidden_dimension = int(hidden_dimension) |
| self.algorithm_count = int(algorithm_count) |
| if feature_mode not in {"token_summary", "structural_lengths"}: |
| raise ValueError("unknown router feature_mode") |
| self.feature_mode = str(feature_mode) |
| self.digit_embedding = nn.Embedding(16, self.digit_dimension) |
| self.role_embedding = nn.Embedding(3, self.digit_dimension) |
| self.length_embedding = nn.Embedding( |
| self.maximum_digits + 1, self.length_dimension |
| ) |
| per_argument = 3 * self.digit_dimension + self.length_dimension |
| network_input = 3 * per_argument if feature_mode == "token_summary" else 9 |
| self.network = nn.Sequential( |
| nn.Linear(network_input, self.hidden_dimension), |
| nn.GELU(), |
| nn.Linear(self.hidden_dimension, self.algorithm_count), |
| ) |
| self.config = { |
| "maximum_digits": self.maximum_digits, |
| "digit_dimension": self.digit_dimension, |
| "length_dimension": self.length_dimension, |
| "hidden_dimension": self.hidden_dimension, |
| "algorithm_count": self.algorithm_count, |
| "feature_mode": self.feature_mode, |
| } |
|
|
| def _row_lengths(self, rows: Sequence[Sequence[int]]) -> torch.Tensor: |
| if not rows: |
| raise ValueError("router arguments must be nonempty") |
| device = self.digit_embedding.weight.device |
| lengths = torch.tensor([len(row) for row in rows], device=device) |
| if bool((lengths < 1).any().item()): |
| raise ValueError("router argument rows must be nonempty") |
| if bool((lengths > self.maximum_digits).any().item()): |
| raise ValueError("router argument exceeds maximum_digits") |
| for row in rows: |
| if any(digit < 0 or digit > 15 for digit in row): |
| raise ValueError("router digits must be in [0, 15]") |
| return lengths |
|
|
| def _argument_features( |
| self, rows: Sequence[Sequence[int]], role: int |
| ) -> torch.Tensor: |
| device = self.digit_embedding.weight.device |
| lengths = self._row_lengths(rows) |
| width = int(lengths.max().item()) |
| digits = torch.zeros((len(rows), width), dtype=torch.long, device=device) |
| mask = torch.zeros((len(rows), width), dtype=torch.bool, device=device) |
| for index, row in enumerate(rows): |
| values = torch.tensor(row, dtype=torch.long, device=device) |
| if bool(((values < 0) | (values > 15)).any().item()): |
| raise ValueError("router digits must be in [0, 15]") |
| digits[index, : len(row)] = values |
| mask[index, : len(row)] = True |
| role_ids = torch.full_like(digits, int(role)) |
| embedded = self.digit_embedding(digits) + self.role_embedding(role_ids) |
| mean = (embedded * mask.unsqueeze(-1)).sum(1) / lengths.unsqueeze(-1) |
| first = embedded[:, 0] |
| last_indices = (lengths - 1).view(-1, 1, 1).expand( |
| -1, 1, self.digit_dimension |
| ) |
| last = embedded.gather(1, last_indices).squeeze(1) |
| length = self.length_embedding(lengths) |
| return torch.cat((mean, first, last, length), -1) |
|
|
| def _structural_length_features( |
| self, |
| a_rows: Sequence[Sequence[int]], |
| b_rows: Sequence[Sequence[int]], |
| modulus_rows: Sequence[Sequence[int]], |
| ) -> torch.Tensor: |
| lengths = torch.stack( |
| ( |
| self._row_lengths(a_rows), |
| self._row_lengths(b_rows), |
| self._row_lengths(modulus_rows), |
| ), |
| -1, |
| ).float() |
| linear = lengths / float(self.maximum_digits) |
| logarithmic = torch.log2(lengths) / torch.log2( |
| torch.tensor( |
| float(self.maximum_digits), |
| device=lengths.device, |
| ) |
| ).clamp_min(1.0) |
| operand_minimum = logarithmic[:, :2].min(-1).values |
| operand_maximum = logarithmic[:, :2].max(-1).values |
| modulus = logarithmic[:, 2] |
| minimum_is_shorter = lengths[:, :2].min(-1).values.lt(lengths[:, 2]) |
| return torch.cat( |
| ( |
| linear, |
| logarithmic, |
| (operand_minimum - modulus).unsqueeze(-1), |
| (operand_maximum - modulus).unsqueeze(-1), |
| minimum_is_shorter.float().unsqueeze(-1), |
| ), |
| -1, |
| ) |
|
|
| def forward( |
| self, |
| a_rows: Sequence[Sequence[int]], |
| b_rows: Sequence[Sequence[int]], |
| modulus_rows: Sequence[Sequence[int]], |
| ) -> torch.Tensor: |
| batch = len(a_rows) |
| if not batch or len(b_rows) != batch or len(modulus_rows) != batch: |
| raise ValueError("router raw arguments must be nonempty and aligned") |
| if self.feature_mode == "structural_lengths": |
| features = self._structural_length_features( |
| a_rows, b_rows, modulus_rows |
| ) |
| else: |
| features = torch.cat( |
| ( |
| self._argument_features(a_rows, 0), |
| self._argument_features(b_rows, 1), |
| self._argument_features(modulus_rows, 2), |
| ), |
| -1, |
| ) |
| return self.network(features) |
|
|
| def route_ids( |
| self, |
| a_rows: Sequence[Sequence[int]], |
| b_rows: Sequence[Sequence[int]], |
| modulus_rows: Sequence[Sequence[int]], |
| ) -> torch.Tensor: |
| return self.forward(a_rows, b_rows, modulus_rows).argmax(-1) |
|
|
|
|
| @dataclass(frozen=True) |
| class AutoRoutedU3Prediction: |
| answers: torch.Tensor |
| algorithm_ids: torch.Tensor |
| executor: UniversalU1Prediction |
|
|
|
|
| class AutoRoutedUniversalU3(nn.Module): |
| """One checkpoint combining the learned router and frozen U2 executor.""" |
|
|
| def __init__( |
| self, |
| *, |
| executor: UniversalSelfClockedU2, |
| router: RawInputAlgorithmRouter, |
| ) -> None: |
| super().__init__() |
| self.executor = executor |
| self.router = router |
| self.config = { |
| "executor": dict(executor.config), |
| "router": dict(router.config), |
| } |
|
|
| def predict_encoded( |
| self, |
| a_rows: Sequence[Sequence[int]], |
| b_rows: Sequence[Sequence[int]], |
| modulus_rows: Sequence[Sequence[int]], |
| ) -> AutoRoutedU3Prediction: |
| algorithm_ids = self.router.route_ids(a_rows, b_rows, modulus_rows) |
| execution = self.executor.predict_encoded_ids( |
| a_rows, |
| b_rows, |
| modulus_rows, |
| algorithm_ids.detach().cpu().tolist(), |
| ) |
| return AutoRoutedU3Prediction( |
| execution.answers, |
| algorithm_ids, |
| execution, |
| ) |
|
|
| def forward( |
| self, |
| a_rows: Sequence[Sequence[int]], |
| b_rows: Sequence[Sequence[int]], |
| modulus_rows: Sequence[Sequence[int]], |
| ) -> AutoRoutedU3Prediction: |
| return self.predict_encoded(a_rows, b_rows, modulus_rows) |
|
|
|
|
| __all__ = [ |
| "AutoRoutedU3Prediction", |
| "AutoRoutedUniversalU3", |
| "RawInputAlgorithmRouter", |
| ] |
|
|
