"""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", ]